Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models
Authors: Lior Cohen
Ofir Nabati
Kaixin Wang
Navdeep Kumar
Shie Mannor
Venue: ICLR 2026 Poster (2026)
Keywords: world models
diffusion
model-basedreinforcement learning
Primary area: reinforcement learning

Abstract: We study diffusion-based world models for reinforcement learning, which offer high generative fidelity but face critical efficiency challenges in control. 
Current methods either require heavyweight models at inference or rely on highly sequential imagination, both of which impose prohibitive computational costs. 
We propose Horizon Imagination (HI), an on-policy imagination process for discrete stochastic policies that denoises multiple future observations in parallel. HI incorporates a stabilization mechanism and a novel sampling schedule that decouples the denoising budget from the effective horizon over which denoising is applied while also supporting sub-frame budgets.
Experiments on Atari 100K and Craftium show that our approach maintains control performance with a sub-frame budget of half the denoising steps and achieves superior generation quality under varied schedules. 
Code is available at https://github.com/leor-c/horizon-imagination.

Published as a conference paper at ICLR 2026
HORIZON IMAGINATION: EFFICIENT ON-POLICY ROLLOUT IN DIFFUSION WORLD MODELS
Lior Cohen' Technion
Ofir Nabati Technion
Kaixin Wang Microsoft Research
Navdeep Kumar Technion
Shie Mannor Technion
ABSTRACT
We study diffusion-based world models for reinforcement learning, which offer high generative fidelity but face critical efficiency challenges in control. Cur- rent methods either require heavyweight models at inference or rely on highly sequential imagination, both of which impose prohibitive computational costs. We propose Horizon Imagination (HI), an on-policy imagination process for dis- crete stochastic policies that denoises multiple future observations parallel. HI incorporates a stabilization mechanism and a novel sampling schedule that de- couples the denoising budget from the effective horizon over which denoising is applied while also supporting sub-frame budgets. Experiments Atari 100K and Craftium show that our approach maintains control performance with a sub-frame budget of half the denoising steps and achieves superior generation quality der varied schedules. Code is available at htttps:: //t::/gitttt cww/ comm/eorr-an c/ horizon imagiimmaaiin nat. iatiagiitiioo
INTRODUCTION
World models (Ha & Schmidhuber 2018) have emerged as a powerful paradigm for achieving sample-efficient reinforcement learning (RL). By learning a generative model of environment dy- namics, world models enable agents to produce large quantities of simulated experience that can be used train controllers, thereby reducing reliance costly real environment interactions. Recently, diffusion world models have become particularly attractive, owing to their unmatched generative fi- delity across modalities such images and video (Esser et al.. 2024 Blattmann al. 2023).
Despite these advantages, current approaches face key limitations when applied to control. Practi- cal deployment often requires controllers that are lightweight and highly computationally efficient, enabling real-time inference with low power consumption. In contrast, the prevailing trend in world models, particularly diffusion- based models, is toward ever-larger architectures (Agarwal et al.) 2025 Parker-Holder et al. 2024: Ball et al. 2025 Assran al.) 2025). This poses significant bar- rier to their direct use at test time as generation becomes increasingly time- and compute-intensin Furthermore, methods that rely repeated interaction with the world model during controller in- ference are therefore impractical in many real-world settings.
Alternatively, approaches that train independent controller via world model imagination (Alonso ct al. 2024) encounter prohibitive computational overhead with diffusion- baasd models. Since each observation must be generated through costly multi-step denoising process, and imagination quires sequentially interleaving policy decisions with world model predictions, the resulting imagi- nation process is inherently sequential and computationally intensive.
In this work, we propose a more efficient approach for training discrete stochastic policies with diffusion world models. We introduce Horizon Imagination, an on-policy imagination process that denoises multiple future observations simultancously, reducing the sequential burden of diffusion generation. We devise (i) mechanism to mitigate policy-induced instability during horizon imagi- nation, and (ii) novel sampling schedule that disentangles the denoising budget from the schedule's decay horizon, enabling independent control over both parameters, supporting sub-frame budgets, and yielding superior generation quality at higher budgets. Our approach is training- and applies to any pre-trained world model with observation-level time conditioning.
"Email: liorcohenS @ @ww.@ac..acccmiiiiii.. .cc.iio..ccciii.c.i.........

---

Published as a conference paper at ICLR 2026
We validate our approach subset of Atari 100K and Craftium environments, showing that the agent maintains control performance with only half the denoising budget. We further analyze world model generation quality under different configurations of the proposed Horizon schedule, spanning fully autoregressive to highly parallel regimes and a range of denoising budgets. Parallel generation consistently proves advantageous, achieving strong performance even under sub-frame budgets.
RELATED WORK
Large Diffusion World Models Diffusion-basee methods have emerged as the state of the art in generative modeling, delivering unmatched fidelity across modalities such as images and video (Ho et al. 2020; Rombach et al. 2022, Ho et al., 2022). Building on this success, recent large-scale
world models adopt diffusion backbones for their superior generative capabilities (Ball al. 2025 Agarwal et al. [025 Decart et al.
2024 [Assran et al. 2025). While motivated by applications such
as agent training, these works are limited conditional video generation and do not address control.
Concurrent Multi-step Generation The idea of generating multiple steps in parallel with diffu- sion world models has been explored in several recent works (Chen et al., 2024: Rigter et al. 2024] Ding et al. 2024, Jackkson et al., 2024). In partieular, Diffusion Forcing (Chen al. 2024 intro- duces multi-step simultaneous generation scheme together with a control framework in which the world model serves as both policy and planner, akin to model predictive control (Gareía et al. 1989). Although effective for learning control, this approach is computationally intensive and often imprac- tical for deployment. contrast, our work addresses this limitation by enabling efficient training in imagination, allowing lightweight policy to be extracted from the world model for deployment.
Rigter et al. (024) and Jackson et al. (2024) proposed policy-guided imagination methods where during the denoising process the actions updated in the direction of the score of the policy action distribution (a)(s as in Langevin dynamics (Song & Ermon 2019) and classifier guidance
(Dhariwal & Nichol 2021). However, these approaches are limited to continuous action spaces.
Ding et al. (2024) proposed an off- off-policy multi-step generation approach for the offiin RL setting. Importantly, none of the above approaches examined how world model generation quality is affected by sampling configurations, including sequential versus parallel generation and the influence of the generation budget.
Diffusion World Model Agents for Control Alonso al. (2024) and Yang et al. (2024) proposed diffusiio-basse world model methods that follow a traditional step-by-sten sequential imagination. Both reported significant computational overhead. In particular, Alonso et al.] (2024) provided detailed runtime analysis, showing that sequential imagination dominates the overall cost.
PRELIMINARIES
Reinforcement Learning Setup Consider a Partially Observable Markov Decision Process (POMDP) environment. Here, we consider practical agnte-agn formulation, as the agent has no knowledge about the POMDP hidden state space. some initial observation 01, the agent picks an action a, and the environment evolves according to Ot+1; rt, rt ~ p(o Ft, de (0gi;aç;) where r, E R.,, de E {0, 1} are the reward and termina- tion signals, respectively. This process repeats until a non-zero termination signal is observed. The agent's 's goal is to maximize its expected return E E 0 rt...
each step
1,2,.... starting from
World Model Agents In reinforcement learning (RL), "world models" (Ha Schmidhuber 2018) denote class of model-based methods where the controller is trained entirely simulated rollouts generated by learned dynamics model, whereas real environment interactions are used for training the model. World-model agents typically consist of four components: representation model for encoding and decoding observations, a dynamics model, a controller, and a replay buffer.
World model agents trained either in online components jointly from scratch, fitting the representation and dynamics to real interactions while optimizing the controller via model-genene gener-te rollouts, whereas offline training uses pre-collected
offline setting. Online training learns the
2

---

Published as a conference paper at ICLR 2026
Ground Tuth
Drsund-rum
mith
Durs
DUTS
OUTS
CUIS
OUT
CD's
Figure I: Example of generation instabilities observed in Crafti um/ ChepTree-yy vO under naive action sampling during horizon imagination. In contrast, our stable sampling method produces robust, high-quality generations. The first context frame is highlighted with a blue border.
data to train the components sequentially: representation, dynamics, and controller. Our method is applicable in both settings.
Diffusion Framework In this work, we employ the I-Rectified Flow (RF) framework (Liu et al, 2023 for its simplicity and robustness, though our method remains compatible with other diffusion variants. Formally, given samples from an unknown target distribution Pdata RF learns time- dependent vector field vo parameterized by 9 that transports samples x. Pprior from a chosen prior distribution Pprior toward data samples x² p² Pdata- The evolution is governed by
dx = v0("", ,T) .T)dT dr, T TE [0,
where superscripts denote denoising time. The vector field ue is trained to estimate E(x. x0 |x], with -xx + (1 T)x", by minimizing the regression objective
((0) Ex ||10 .T (x* - -00)1²
4 METHOD
Multi-step trajectory generation naturally alternates between producing an observation o, and sam- pling an action at x(- Ogt, act), as each step depends on the output of the previous one. How- ever, since the generative process of diffusion methods involves multiple (costly) forward passes for generating each observation, the above sequential generation of trajectory becomes highly sequential and expensive, both in time and in compute.
We propose a more efficient approach where the denoiser vy denoises multiple observations in paral- lel, conditioned on actions that are updated to track the policy as it co-evolves with the world model and incorporates newly available information. However, because action changes alter subsequent observations through different dynamics, they may in turn trigger further updates to both the pol- icy and observations and destabilize the denoising process (Figure [].. Our approach therefore aims to minimize unnecessary action changes due to stochasticity, while ensuring that actions remain aligned with the evolving policy.
4.1 WORLD MODEL TRAINING
The world model a learned generative model of the environment dynamics
p(z++1, rty de Zgge, age),
operating over sequences of compact latent representations of observations These la- tents are produced by the encoder of the representation model, whose output is constrained to z: de via tanh activation. The world model comprises denoiser ty for gen- erating future observations and a separate, lightweight terainntttnnnnoon predictor. This design cleanly separates large, general-purpose dynamics module from small, task-specific predictors.
3

---

Published as a conference paper at ICLR 2026
During training, h-step trajectory segments comprising observations O1:h, actions a1:hs rewards, and terminations are sampled from the replay buffer. The observations 01:h are encoded into latent representations Zi:h, serving as the clean targets z' Pdata- A uniform noise sample z Ppeior = U([-1, 1))xxdexxii Xdux of matching shape is also drawn. Following Chen et al. 2024), an independent denoising time is sampled for each observation, i.e., 5 ~ (4 (1,00))) With probability 0.2, a clean prefix is provided by setting 7sk where k ~ U([I,...... [0.7h])) We empirically found that it improves generation quality by better matching inference-time conditions, where the initill context noise-free.
Given z'
Pdaaa. z"
Pprior- and the denoiser vy is trained to approximate
Exp xt [z]
z!. z
acc rec.
AA Vists h.
by minimizing the rectified flow regression objective
h Ex 21 (ll0
L(O)
27 Tc)) - (z) =z,)1²
(1)
h 1=1
where 2 Ti 2 + (1-1)) TB) z.. All denoiser outputs are computed in parallel in a single forward pass. Each output corresponds to particular observation at timestep and uses only inputs up timestep, preserving causality. Further details, including teranaiternination modeling, are provided in Appendix 1.2 Pseudocode is given in Appendix C
4.2 HORIZON IMAGINATION (WORLD MODEL INFERENCE)
Denoising multiple observations in parallel during imagination implies denoising far-future obser- vations before near-future ones are fully clean. Moreover, denoising such far-future observations requires access future actions up to that timestep (see Eq. [ These actions therefore must be conditioned on intermediate, still-noisy observations. We therefore query the policy before each denoising step, yielding an increasingly informed categorical action distribution for each timestep. Later, the policy is optimized across all noise levels based on its imagined interactions.
However, naively sampling from these distributions can lead frequent action changes across de- noising steps, particularly when the policy outputs high-entropy distributions, even when the policy itself remains fixed. To address this, we introduce a stable sampling method that reduces unneces- sary action changes throughout the denoising process. For clarity, we present the method at the level of single action some arbitrary time
Minimizing Action Changes Over the Denoising Process Consider set of N discrete actions
A = [N]. At denoising time т, the policy induces categorical distribution x π x- acr) E A N-1 where N-1 is the standard (N 1)-simplex. Inspired by inverse transform sampling, method leverages a single multivariate uniform sample that is consistently mapped to an action under the evolving distributions T..
At the beginning of the denoising process, we draw two random objects: (i) a vector w ~
u((0, 1), and (ii) random permutation p [N] - [N] that fixes an order over the actions. At each denoising time r. the action is given by a = a( wW), ensuring that the same uniform sample w yields consistent decisions across the evolving distributions x- The action a(TT, w) is determined by scanning w's coordinates:
((i)
for the smallest i with wi < (z(T)),
a(TTWw) ()
(2)
((N) no such exists,
where
Kettt S(
S,(T) > S,(T) 0
N
(1((T)
S.(T) = Σ -P()
Tp(()
j=1
Intuitively, (=) the conditional probability of choosing the i-th action in the order p. given that all earlier actions were skipped.
4

---

Published as a conference paper at ICLR 2026
Denaising steps
Frame index
15
Horizon schedule
1.0
0.6
Pyramidal schedule
0.4
0.2
budget=10
budget=20 20
Figure 2: Comparison of the Pyramidal schedule (Chen et al. 2024 and the proposed Horizon
schedule (transposed). Horizon fixes the decay horizon (v = 3), yielding consistent schedules across budgets, whereas in the Pyramidal schedule the decay horizon drifts with budget, the two are entangled, leading to degraded generation quality at higher budgets.
Proposition The sampling scheme a(..-) satisfies the following properties:
1. For any distribution p A N-1, the random variable a(p,w) with u([0, N-) is distributed according to P, Le., -)p,-)
~ -)
2. For any P. N-1 (e.g., p Th and T2 with TI T2), consider the event
A {w a(p,w) # w))
Let 5ct 9) = 1pp 4/l2 denote the total variation distance between p and q. Then
8(p.q) q) Pr(A) I ||a(p) - aag))l
We defer the proof to Appendix E. As an immediate corollary, if p y then Pr(A) O. meaning that when the distribution does not vary between denoising steps, the actions necessarily remain unchanged. In contrast, naively sampling fresh action at each step is likely induce frequent action changes under sufficiently high-entropy policies.
The Horizon Schedule We propose a novel time schedule inspired by the pyramidal schedule of Chen et al. (2024), which denoises near future observations before far future ones. Unlike prior approaches, our schedule disentangles the rate at which denoising progresses across time, i.c., the decay horizon, from the denoising budget, i.c., how many denoising steps are performed in total. This design provides precise and independent control over both parameters.
Formally, let [1, h] specify the decay horizon, i.c., the number of steps which the denoising schedule decays from to O, and let B N denote the total budget of denoising steps. We define the sampling schedule as matrix K E [o. 1] ()) where each of the first rows specifies the denoising times assigned to all h observations at given step. The last row ensures that 1 for all < t < h. To determine the values of the entries of K, consider the lines
b h-
b)),b)
-11
with slope . sequence time index 0 S < h, and denoising step index 0 < b S B. The entries of K are given by
Ki, clamp(<(j 1, 1),0, 1).
(3)
5

---

Published as a conference paper at ICLR 2026
Algorithm Horizon Imagination
Require: World model denoiser policy generation horizon h, denoising budget B, decay horizon v, reward and termination models rab dæ, optional context segment (Zgk, ack) with
k 620 0
1: Compute the Horizon schedule K given h, B,
EEq.
2: Sample 2++1" "K11..... 4
Pprior
3: Sample Wt, Pt for every k + 1sss k h
4: if k >> then
Sample first action ak ~ T(- Zck. ack) 6: end if
5:
7: for each denoising iteration b from to B do
8: Set < K5 and dr Kb+1: Ko::
9: Denote 22444 (z 1+11 xxt
10: for ktt do
11:
Compute action distributions Txt t x(- 2 sate a akkto: Generate actions air. = (TK+e:
12:
(Eq. 2b
13:
end for
14:
Perform denoising step: Zk+t
trr +- zkut + vo Uo a aleteer Tg)) dr: (single forward pass)
for all 1 < < h with dr, 0
end for
16:
Compute rewards and terminations Eskth, dekth with ro: de
17: return imagined trajectory 2skth= ackth; Fgkth, dgath
Importantly, the proposed schedule allows any combination of B and v, and particular h. Thus, while standard auto-regressive generation requires B 2 h, and is limited to multiples of h, our schedule can use arbitrary positive integer budget values, including sub-frame budgets.
4.3 ACTOR-CRITTE TRAINING
In the previous sections, we described how the interaction between the policy and the world model unfolds. These interactions generate the data used to optimize the controller. Unlike standard imag- ination, our parallel generation requires the policy to produce outputs at every denoising step, from pure noise fully denoised samples. Consequently, the policy must be trained across all noise levels to ensure meaningful outputs.
We build on an actor-critic method introduced in prior works (Cohen et 2025 Hafner et al., 2025). The actor and critic are implemented as separate networks with identical backbones and distinct linear heads for their respective outputs. The critic operates only on fully denoised inputs, providing bootstrapped value estimates for the actor. The actor is trained using the standard REIN- FORCE objective (Sutton al. 1999) applied at every trajectory step before termination, leveraging interactions generated across all noise levels. Formally, conditioned on context ziga ack of & clean frames, the policy objective for trajectory time k + and denoising time is given by
AR++ log x (a 2k41 t acktt. + m
where A, is the advantage at step /. and H. is the policy entropy. To ensure balanced policy updates over denoising times, we restrict updates to steps which the denoising time of the next observation increases. Full actor-critic details given in Appendix A.1.3
EXPERIMENTS
5.1 EMPIRICAL ANALYSIS ACTION CONSISTENCY
We conduct two controlled experiments to evaluate the proposed stable action sampling mechanism.
In the first experiment (Figure Baj, we examine the empirical distribution average action changes between random source and target distributions p, q using method. We sample 10d pairs of

---

Published as a conference paper at ICLR 2026
Stable (ours)
Naive
Number of action changes
2.0
Average number of changes
12.3
10.0
1.0-
1.4
5.0
0.0
64
balanced
high
Number Actions (N)
Entropy Setting
(a) Distributions of the average number of action changes, 6(p q) units, for various values. Minimum, mean, and maximum are indicated.
(h) Comparing the number of action changes be- tween naive sampling and the proposed method (mean + std).
Figure 3: Empirical study of the average number action changes under various settings.
categorical distributions from a uniform Dirichlet prior with N actions. For each pair, we draw 10 (w, p) samples, compute the induced actions, and record the average number of changes. This repeated for various N values. Empirically, our method performs very close to the tight the- oretical lower bound, the total variation 8(p, q), across all N values, demonstrating near-optimal performance.
compare naive sampling with the proposed stable mech- anism a simulated denoising process. We interpolate between source and target categorical distribution with N = 10 actions across 8 steps, then fix the target distribution for additional steps. Both source and target are sampled from a Dirichlet distribution under three entropy settings: low (0.2. ... 0.2), uniform, and high (5,... ... 5). For each entropy setting, simulate the process for 1,000 source-target pairs and estimate the average number of action changes across all 16 steps 106 simulations per pair. Our results show that while the proposed method yields most action change on average across all 16 steps, naive sampling incurs changes in most steps, altering more than half overall. For high-entropy distributions, the number of action changes decreases under our method but increases with naive sampling, consistent with our theoretical results.
In the second experiment (Figure 3b),
5.2 CONTROL PERFORMANCE
To study whether Horizon Imagination enables efficient training, benchmark our world model agent in an online RL setting. The agent's training process involves a repeated cycle of four stages: data collection through environment interactions. representation learning, world model learning, and actor-critic learning in imagination. For comparison, we consider three baselines that differ only in their imagination configuration: standard autoregressive baseline with 32 denoising steps in total (one step observation), denoted (v 1, B = 32), and two variants with decay horizon v 4 and budgets of 16 and B
5.2.1
SETUP
Agent Implementation For latent representations, we employ the contimuous autoencoder image tokenizer of Agarwal et 2025)). and apply final tanh activation the encoder output con- strain values to 11 The denoiser network up is implemented as an action-conditions causal Diffusion Transformer (DiT) (Peebles & Xie 2023) architecture. This model operates on sequenees of latent observations, where action and denoising time conditioning are carried via the adaptive layer normalization (AdaLN) layers of the DiT architecture. Rewards and termination signals are modeled using a lightweight recurrent neural network (RNN) with compact convolutional neural network (CNN) feature extractor and two output heads, operating on encoded trajectories. The same RNN architecture is used for both the actor and critic networks. For simplicity, omit action conditioning in the actor, critic, and the observations alone provide sufficient information. The agent uses fixed hyperparameters across
models, as in all benchmarks considered

---

Published as a conference paper at ICLR 2026
all experiments and comprises tokenizer (22.5M), world model (67M), and actor-critic (7.5M), totaling 97M parameters. We defer the full implementation details to Appendix
Benchmarks We evaluate on four Atari100K (Kaiser et al. 2020 games and four Craftium (Malagón et al., 2025) games, both benchmarks involving visual inputs and discrete actions. All environments were run for 100K interaction steps, except Craft ieffiium ium SmallRoom-vo which was capped at 30K due to its lower difficulty. The well-established Atari benchmark serves as solid control evaluation suite and is visually simpler. In practice, we found that even a single de- noising step can yield satisfactory generation quality in Atari. By contrast, Craftium presents richer and more complex observations, providing a more challenging test for the world model.
conducted NVIDIA A100 GPUs, except for Craftiun um/ Small Room-vo, which was run on RTX 5090, RTX 4090, and A40 GPUs. On Atari, training with denoising budget of B 32 required approximately 27 hours per run, while reducing the budget to B 16 shortened training to about 19 hours. Because the tokenizer and world model training stages are unaffected by this parameter, the overall speed-up is smaller than full 2x. On Craftium, runs take longer due to slower environment interactions.
Training Time and Hardware All experiments
5.2.2 RESULTS
As shown in Figure both baselines with decay horizon maintain the performance of the autoregressive baseline across all environments, achieving comparable results at reduced cost. No- tably, our approach sustains full performance with a sub-frame budget of 16, requiring only half the denoising steps of the autoregressive baseline.
Moreover, our results suggest that single denoising step per observation generally suffices, with baselines at B = 32 performing comparably across environments. In contrast, in Cra ft um/- ChopTree-vO, the most visually complex environment, the (v 4, B 32) baseline achieves superior performance, suggesting improved effectiveness under the B 32 budget.
(v 1, 32) (AR)
16)
32)
Craflium/Smilllemm--o
Craftiam/loomman
Craftamuppslleeuu
Craftium/ChnTTre-p
-200
Epionde return
14-
Enlodd eatumn
+400
20
-320
1.
-420
-030
20
-530
Enuironment Steps
Environment Steps
Envilonmmnt Steps
Envvirommnt Steps
WLExBoxing
ABBirasyCimberrs
ALEGoopherUS
2202
KOL
Eplonde reTUET
Epicadie return
-55
42
Epicode
Eploode
Z0k
12kc
60k
20k
25c
50c
100k
22k
75k
1006
20h
406
BOK
120K
25k
526
736
200k
Environment Steps
Environment Stegs
Environment Stops
Environment Steps
Figure Actor-Critic Performance. Average episodic return curves of key baselines during training. Each baseline is evaluated over 5 seeds. Curves show the mean and standard deviation, smoothed by a moving average (window size 15). dashed horizontal line denotes Atari human-level perfor- mance.
8

---

Published as a conference paper at ICLR 2026
5.2.3 ABLATION STUDY: STABLE ACTION SAMPLING
To assess the influence of the proposed stable action sampling mechanism on control performance, we perform an ablation study by reverting to naive action sampling. In this baseline, fresh action is independently drawn from the policy distribution before each denoising step.
The results in Figure 5 show that replacing the stable action sampling mechanism with naive sam- pling leads to substantial drop in control performance (returns). Notably, while the naive baseline consistently weaker, we observe a particularly prominent collapse on Atari Boxing and Gopher. These findings further highlight the importance of preventing unnecessary action fluctuations during the denoising process.
Stable (v 4, B 16)
Naive (1/ 4, 16)
Craftuum/mallRcom-000
CraftiumRRoomw0
Cratkum/Spelec-v0
Craftum/ChosTeee-g
20
-120
Episande return
Faioade ratur
12
-250
M
Egicada
M
-2k
37k
232
522
732
100c
73t
170k
27a
40c
80.
2004
Ensironment Stegs
Environment Steps
Ertril cevirnttnnntt Steps
Environment Steps
ALE-BBoxinggs
A.L:GrayCCimbeen
ArErGepheruS
ALE/Runggi uNaster-d
02-
Sc
02
Episode returrn
Eploode return
-25c
Epiicde
Episode
25kc
J0k
25c
50c
75
100k
23k 50k
75k 100e
25k
402
60c.
BOk 100k
25c Environment Steps
52t
75k
200k
Environment Steps
Environment Steps
Enivtromment
Figure 5: Actor-critic performance comparison between the proposed stable action sampling method and the naive baseline. Each baseline evaluated over seeds. Curves show the mean and standard deviation, smoothed by moving average (window size 15). A dashed horizontal line denotes Atari human-level performance.
5.3 PARALLEL VS. SEQUENTIAL GENERATION QUALITY
To study the impact of parallel multi-step denoising generation quality, evaluate the trained world models from Section 5.2 under varying decay horizon (1/ and denoising budget (B) values. Our central question is how generation quality differs between parallel and sequential configurations, and which settings prove most effective.
In our evaluation, 512 episode segments of 33 frames sampled from the training data. For each combination of and B. 512 corresponding 32-frame continuations are generated using the first frame as context. To isolate the impact of and B from other factors such the action sampling scheme, the original recorded actions are provided to the model at every denoising step. Generation quality is assessed using Fréchet Video Distance (FVD) (Unterthiner al. 2019) and mean squared error (MSE) against the ground-truth segments. FVD better reflects perceptual quality, while MSE more indicative deviations from the ground truth sequence.
show a consistent FVD trend across environments, indicating a clear advan- tage for parallel configurations (4 S V < 16) under low to medium budgets. Remarkably, in several cases, parallel variants with sub-frame budgets achieve quality comparable to baselines that rely the maximal budget, 16 times larger. At extreme budgets (B 2 128), performance tends to degrade slightly as increases, with the autoregressive baseline consistently ranking at the top. Interestingly,
Our results (Figure 6

---

Published as a conference paper at ICLR 2026
Decay
(AR)
16
32
Craftuu/Spolko-vo
To
Craftiun/Reem-v
SentuumSmalReom--
145
122
00
102
300
03-
60
62
200
16
126
16
60
40
120
128
Denaising Budget (8)
Denoising Budget (9)
Denoising Budget (R)
Denoising Budget (B)
Craftuum/Speloc-v0
Crattum/ChapTee--0
Crafflum.Rcom-vo
Cottlum/SmalRoomm.v
0.0220
Main Squared
0.0316
0.014
0.015
0.00
0.013
6.34
0.03
0.014
0.0312
35
128 200
36
32
04
322
250
36
52
24
123
230
15
66
128 255
Denolsing Budget (B)
Denolsing Budget
Denolsing Budget
Danolsing Budget 18)
Figure World model generation quality versus denoising steps budget. Each point shows the average FVD/MSE over 512 sampled 33 33-frame segments, where the first frame was given as context and the last were generated conditioned on the recorded actions. A dashed vertical line indicates the transition out of sub-frame budgets.
the MSE results (bottom row of Figure 6 suggest that increasing the denoising budget may cause generated sequences to drift further from the ground truth, even perceptual quality improves.
Finally, in Appendix D we present the corresponding Atari results and further repeat this experi- ment under the Pyramidal schedule of Chen al. (2024), which suffers pronounced collapse in performance as budgets grow.
6 LIMITATIONS
While our control results consistently highlight the advantage of v 4, we did not extend our study additional configurations. Evaluating each baseline requires running full experiments across all environments with five seeds environment, which is computationally prohibitive. Consequently, we had to restrict our study to eight environments keep the evaluations tractable. Future work could explore additional configurations.
7 CONCLUSIONS
In this work, we addressed the computational challenge of on-policy imagination in diffusion world models with discrete actions, motivated by the need to train lightweight policies for deployment. We introduced Horizon Imagination (HI), parallel multi-step imagination procedure that incorpo- rates two novel mechanisms: stable discrete action sampling and the Horizon schedule. Our analy- sis shows that stable sampling ensures consistent action selection, thereby stabilizing imagination, while the widely applicable Horizon schedule delivers sub-frame performance and preserves robust decay horizons across budgets. Together, these advances allow HI sustain strong performance even under sub-frame budgets and at half the computational cost, making diffusion world models markedly more practical for training deployable policies.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Niket Agarwal, Arslan Ali, Maciej Bala, Yogesh Balaji, Erik Barker, Tiffany Cai, Prithvijit Chat- topadhyay, Yongxin Chen, Yin Cui, Yifan Ding, et al. Cosmos world foundation model platform for physical ai. arXiv preprint arXiv:2501 250.:25000.03355, 2025.
David Aldous. Random walks on finite groups and rapidly mixing markov chains. In Jacques Azéma and Mare Yor (eds.), Séminaire de Probabilités XVII 1981/82, pp. 243-297, Berlin, Heidelberg, 1983. Springer Berlin Heidelberg. ISBN 978-3-540-33614-77
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. In A. Glober- son, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 58757-58791. Curran Associates, Inc., ngs. www.wps..pp.... wttp:ppppppppppeceee fttpp/errrileeepppe////00 2024. URL F12200eeeeTTT1133332222320000303 coneeeeee pdf
Mahmoud Assran, Adrien Bardes, David Fan, Quentin Garrido, Russell Howes, Mojtaba Komeili, Matthew Muckley, Ammar Rizvi, Claire Roberts, Koustuv Sinha, Artem Zholus, Sergio Ar- naud, Abha Gejji, Ada Martin, Francois Robert Hogan, Daniel Dugas, Piotr Bojanowski, Vasil Khalidov, Patrick Labatut, Francisco Massa, Mare Szafraniec, Kapil Krishnakumar, Yong Li, Xi- aodong Ma, Sarath Chandar, Franziska Meier, Yann LeCun, Michael Rabbat, and Nicolas Ballas. V-jepa 2: Self-supervised video models enable understanding, prediction and planning. arXiv preprint ar arXiv:2506.09985 09985, 2025.
Philip J. Ball, Jakob Bauer, Frank Belletti, Bethanie Brownfield, Ariel Ephrat. Shlomi Fruchter, Agrim Gupta, Kristian Holsheimer, Aleksander Holynski, Jiri Hron, Christos Kaplanis, Marjorie 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 Baetu, Jordi Berbel, David Bridson, Jake Bruce, Gavin Buttimore, Sarah Chakera, Bilva Chandra, 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 Nandwani, 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, Koray 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. 2025.
Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with latent diffusion mod- els. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 22563-22575, June 2023.
Boyuan Chen, Diego Martf Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitz- mann. Diffusion forcing: Next-token prediction meets full-sequence diffusion. In A. Glober- son, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang (eds.), Advances in Neural Information Processing Systems, volume 37, pp. 1-111-2 Curran Associates, Inc., neurips cc wwwpraappppp____ http://aaee_iilleeeaaee/////2 PaapaaaeeePpeee--CCeeeeeeeeeeee
2024. URL
Lior Cohen, Kaixin Wang. Bingyi Kang, and Shie Mannor. Improving token-based world models with parallel observation prediction. In Forty-first International Conference on Muchine Learning, 2024.
Lior Cohen, Kaixin Wang, Bingyi Kang, Uri Gadot, and Shie Mannor. Uncovering untapped poten- tial in sample-efficient world model agents. arXiv preprint arXhv:2502.11537, 2025.
Decart, Julian Quevedo, Quinn McIntyre, Spruce Campbell, Xinlei Chen, and Robert Wachen. Oa- sis: A universe in transformer, 2024. URL httll/ttss::/oasisseesss
thub. io/

---

Published as a conference paper at ICLR 2026
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans image synthesis. In M. Ranzato, A. Beygelzimer. Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.),
Advances in Neural Information Processing Systems, volume 34, pp. 8780-8794. Curran Associates, Inc., 2021. URL https::/ newrrsssneuuuiee neurips ccpppperrrrppppprrrrrrppprpp files/ 2Bpp://0000/1ee//oooo222lll11112
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model: Fu- ture modeling beyond step-by-step rollout for offiin reinforcement learning. arXiv preprint arXiv:2402.03570. 2024.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion En- glish, and Robin Rombach. Scaling rectified flow transformers for high-resolution image syn- thesis. In Forty-first International Conference on Machine Learning, 2024. URL https: //openreview net. forum? (tt.://oooonn/////PP0 d=FPrUhsQJ55
Carlos García, David M. Prett, and Manfred Morari. Model predictive control: Theory and practice survey, Automatica, 25(3):335- 348, 1989. ISSN 0098-1098. doi: https///oi.org/
10.1016/0005- (990))99020000000999000200000200 URL https: //www/ sci encedi com wecssciiecee
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evo- lution. In Advances in Neural Information Processing Systems 31, pp. 2451-
2463. Curran Associates, Inc., 2018. URL https: nips. cc/paper:
https://
worl herllmodels.githhe
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, 640(8055):6647-533 Apr 2025. ISSN 1476-4687, doi: 10. $41586-025-087444 URLhttps:://ii.or///1 1ttp://///1111111111/111100
Dan Hendrycks and Kevin Gimpel. noniinnen and stochastic stccaacc regulaizzes with gaussian error linear units, 2017.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neu- ral Information Processing Systems, volume 33, pp. 6840-6851. Curran Associates, Inc., 2020. URL Ittps::/pproceedinggss neurips............p...........s cc cc/paper Saaaper_fllesspaaeaaaaeeeeeeeeee filess/aper/222c// f F Elle//cbbciecc3335800066611
79c44b wwwp11999aaaaaaaaa.aaaaa
Jonathan Ho. William Chan, Chitwan Saharia, Jay Whang, Ruigi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303. 2022.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Matthew Thomas Jackson, Michael Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, and Jakob Nicolaus Foerster. Policy-guided diffusion. In Reinforcement Learning Conference, 2024.
URL httpps:://peeeeeiie.. net.:iiwwwweeeeeeeemm forrmm 2ideedd=mmllllacacm
E. Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiáski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In International Conference on Learning Representationss 2020. URL https:// openreview www.riiwwwee/eeeee wwte:/forrmmmid=IxCcRe
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the de- sign space of diffusion-based generative models. In S. Koyejo. S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.). Advances in Neural Infor- mation Processing Systems, volume 35, pp. 26565-26577. Curran Associates. Inc., httpaaaare_.ieeeeeapee/e/// 2022. URL URL.htttsss/prooceeiiss..errrrre www.eeeeeddinnn.....i...........
12

---

Published as a conference paper at ICLR 2026
Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Repre- sentations, 2023. /
Mikel Malagón, Josu Ceberio, and Jose A. Lozano. Craftium: Bridging flexibility and efficiency for rich 3d single- and multi-agent environments. In International Conference on Machine Learning (ICML). 2025.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers are sample-efficient world mod- els. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL
net/ net/pdf?
iddvhFulAcbbxbb
Jack Parker-Holder, Philip Ball, Jake Bruce, Vibhavari Dasagi, Kristian Holsheimer, Chris- tos Kaplanis, Alexandre Moufarek, Guy Scully, Jeremy Shar, Jimmy Shi, Stephen Spencer, Jessica Yung, Michael Dennis, Sultan Kenjeyev, Shangbang Long, Vlad Mnih, Harris Chan. Maxime Gazeau, Bonnie Li, Fabio Pardo, Luyu Wang. Lei Zhang, Frederic Besse, Tim Harley, Anna Mitenkova. Jane Wang, Jeff Clune, Demis Hassabis, Raia Hadsell, Adrian Bolton, Satinder Singh, and Tim Rocktäschel. Genie 2: A large-scale foun- dation world model. 2024. URL
htwpa/eeeeidnooogg...oooo gerie-2-a-large- scale foundation- tto.g///l..coooo///////oooocce scover/blog/
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 1195-4205, October 2023.
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions, 2018. URL https://openrevieet net net. forum d-
Mare Rigter, Jun Yamada, and Ingmar Posner. World models via policy-guided trajectory diffu-
sion. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https:
7 //open/eviww.net/// (atrrmmm?dd-Ccc0000he
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Con- ference on Computer Vision and Patlern Recognition (CVPR), pp. 10684-10695, June 2022.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data dis- tribution. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue Buc, E. Fox, and R. Gareett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL httpps:/pprocss//e
files/ www.ris.../ppaaeeeeeee pdd.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputiing 568:127063, 2024. ISSN 0925-2312. doi: http://ttt..too.eeeeeeececceaaaa
127063. URL http::ttt//////t https://www,
sciencedirecte com
Richard S Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement leaming with function approximation. In S. Solla, T. Leen, and K. Müller (eds.), Advances Neural Information Processing Systems, volume 12. MIT Press, 1999. URL. https::/ ngs nww.wnewr....e httpr//oe...eee.eeeeeerrss.o.... wwtppppr.ccppppp.pppppppppppp..e wwt.:fffllss....e/eeem less/pperr/199// URL Fiee/6/82255000eeeeeee333333---3
Thomas Unterthiner, Sjoerd van Steenkiste, Karol Kurach, Raphael Marinier, Marcin Michalski, and Sylvain Gelly. Towards accurate generative models of video: A new metric & challenges, 2019.
0177
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 on Learning Representations, 2024. URL https:
//openrevies wwwp://e//o.ommtmmommmmm dett://forrmmmdddssssxyxaaaaaa
13

---

Published as a conference paper at ICLR 2026
A EXPERIMENTAL SETUP
A.1 MODEL ARCHITECTURES
A.1.1 REPRESENTATION MODEL (TOKENIZER)
We use the continuous image autoencoder of Agarwal et al. (225 the representation model of our agent. We modify the output of the encoder network by adding a final tanh activation to enforce observed
loss of performance due this modification. In addition, we adjust the configuration to fit our settings and computational budget. Further details on the architecture and hyperparameters can be found in Tables and and in our open-source code.
values in [-1, 1]. Empirically,
We use a standard mean squared error (MSE) regression objective for the reconstruction loss instead of the original L1 objective, we empirically found to perform better in our settings. The full optimization objective is the sum of the reconstruction loss and perceptual loss pperoopual (Agarwal
al. 2025 (Eq. 3, p. 15). Formally, the loss given by
Luokenieer
11ó a|12 ppecceptual:
where
C Ppecceptual
ay [ VGG;(o)ll..
the perceptual loss, == &Pee (ennc (o)) the reconstructed observation produced by the autoen- coder, VGG, are the features from the 1-th layer of a pretrained VGG-19 network with layers, and Og is weight associated with the 1-th loss term. We the same hyperparameter values as Agarwal et al. (2025).
The encoder and decoder architectures of the representation model. "Conv2d(a.b.e)" rep- resents a 2d-convolutional layer with kernel size a x a, stride and padding c. "GN" represents GroupNorm operator with 32 groups, Le and learnable per-channel affine parameters. "Upsample" represents a repeat-interleave operation that doubles the spatial dimensions, followed by Conv2d(3,1,1) layer.
Table I:
Module
Output Shape
Module
Output Shape
Decoder
Encoder
Input
16 x
Conv2d(1,
Input
3 x 64 x 64
Conv2d(3, 1, 1)
256 8 8
Conv2d(3,
61 x 64 6-1
ResnetBlock
256 x 8 8
ResnetBlock
x
AttnBlock
256 8 8
ResnetBlock
128 x 64 64
ResnetBlock
256 8
Conv2d(3, 2, 0)
128 x 32 x 32
ResnetBlock
ResnetBlnck lock
256 x 32 32
ResnetBlock
256 8 x 8
ResnetBlock
256 32
ResnetBlock
256 8 8
Conv2d(3, 2, 0)
256 x 16 16
Upsample
256 x 16 16
ResnetBlock
256 x 16 x 16
ResnetBlock
256
16
AttaBlock
256 x 16 16
AttnBlock
256 x 16 x 16
ResnetBlock
256 16 16
ResnetBlock
256 x 16 x 16
AttnBlock
256 x
AttnBlock
256 x 16 x 16
Conv2d(3, 2, 0)
256 8
ResnetBlock
256 x 16 16
ResnetBlock
256 x x
AttnBlock
256 x
16
AttnBlock
256 x 8
Upsample
256 x 32 x 32
ResnetBlock
256 8
ResnetBlock
128 x 32 32
256 x
ResnetBlock
128 x 32 32
Conv2d(3, L)
16 8 8
ResnetBlock
128 x 32 32
Conv2d(1, 0)
16 8 8
Upsample
128 x
64
tanh
16 8
GN
128 x 64 x 64
Conv2d(3, 1. 1)
3 64 x 64
14

---

Published as a conference paper at ICLR 2026
A.1.2 WORLD MODEL
Our world model consists of two components: the denoiser vy and a separate reward-termination predictor network. We view the dynamics model as a large-scale, general module that captures the complexity of the environment, whereas peward-terminann task-specific, and operate on the dynamics outputs. This design disentangles their dependencies: dynamics models can be scaled and pre-trained advance, while downstream tasks require learning predictors upon receiving new task.
predictors should remain lightweight,
only small, efficient
For the DiT-based denoiser vi (Peebles Xie 2023), we build on the Cosmos implementation (Agarwal al. (025) with several adaptations to setting. As in [[Agarawl et al. (2025). inputs are patchified with 2 x spatial patch layer, and a corresponding reverse patchify layer is applied at the output. We also adopt 3D RoPE positional embeddings (Su et al. 2024 Agarwal et al. 2025
but in contrast to Agarwal et al. (2025), we omit any additional positional encodings.
Since no textual conditioning is required, we remove the cross-attention operator, yielding vanilla DiT block structure. Attention is modified with frame-level causal mask: tokens can attend to all other tokens in the same frame as well as to tokens from preceding frames. To condition on actions, we use a learned embedding table mapping discrete actions to vectors, which are summed with the denoising-time embeddings to form the conditioning input to the AdaLN layers.
Note that we only used the architecture implementation of Agarwal et al. (2025). Specifically, our work is based on the rectified flow diffusion framework (Liu et al. 2023). which is different from the EDM Karras et al. 2022) framework used by Agarwal et al. (2025).
The cerrd-termination network processes latent observation sequences Zi, Zg. analogous va- Each latent is first encoded by a compact CNN consisting of single ResNet block with 256 chan- nels, followed by 1x1 convolution that reduces the dimensionality 64 channels. The resulting features are flattened and passed through linear layer with a Sigmoid Linear Unit (SiLU) activa- tion HHendrycks Gimpel (2017); [Ramachandran al.] (218)), yielding a vector representation for each 3D latent observation. Finally, the sequence of vectors is processed by an LSTM (Hochreiter Schmidhuber 1997) followed by two linear output heads that predict rewards and terminations, respectively.
The reward-termination model is trained on the same clean trajectory segments used for optimizing Ug. To accommodate reward signals that vary widely in scale and sparsity, we model rewards in symlog space. Formally, the reward objective is
|| [ age) - symlog(r:) 11?
where r (Zge+1, a denotes the raw prediction and symlog(x) sig)(r) log([ær + 1) is the symlog transform (Hafner et al. 2025). At inference time, the reward estimate is recovered via
A. symcxpp7r(2ss41 ag:)),
symexp(#) sign(x) (exp([>)) 1),
where symexp is the inverse of the symlog transform.
The termination predictor is trained using standard cross-entropy loss.
A.1.3 ACTOR-CRITIC
The critic T/ is trained exclusively on fully denoised outputs via the following regression objective in symlog space:
h Liv Σ
symlog(G.) V (Zgg,
where symlog(() sign(x) log((x) + 1) is the symlog function proposed by Hafner al. 20255 for robust handling of values across scales and G, is the X-return at step t:
Fr + (11- syyexx((V)) + AGG+1) symexp(VV7)
< h, t= h.
G. =
Here, symexp(x:) = (exp(=)(exp((==)) 1) is the symlog inverse, V = (zzge age)). and F is the reward generated by the learned reward model. We highlight that found the above simple
15

---

Published as a conference paper at ICLR 2026
regression objective in symlog space perform robustly across rewards from dense sparse and across reward scales, as can seen in Figure
Following Hafner et al. (2025), scale the advantage terms using exponential moving average (EMA):
G. Cymexx)x)(V)))
A. ag
(5)
max(1,
where sg denotes the stop-gradient operator and
S = EMA(quantile(GG. (,, 95) - quantile (GatleeGG,a 5), 0.005)
the difference between the 95th and 5th return quantiles within the current batch, smoothed by an EMA with coefficient 0.005.
A.2 HYPERPARAMETERS
For the actor-critic control performance evaluation experiments, we use the same agent hyperpa- rameters across all environments and benchmarks. High-level agent training parameters detailed in Table 2 The hyperparameters for the representation model (tokenizer), world model, and actor- critic, are presented in Tables 4 and respectively.
Table Agent hyperparameters.
Description
Value
Number of epochs
500
Data collection steps per epoch
200
Tokenizer optimization steps per epoch
300
Tokenizer train from epoch
10
World model optimization steps per epoch
300
World model train from epoch Actor-critic optimization steps epoch
25
50
Actor-critic train from epoch
40
Table 3: Representation model (tokenizer) hyperparameters.
Description
Value
Optimizer
AdamW
Learning rate
2e-4
AdamW weight decay
0,05
AdamW (81, B2)
(0.9, 0.95)
Batch size
32
Max. grad norm
Patch size
16

---

Published as a conference paper at ICLR 2026
Table 4: World model hyperparameters.
Description
Value
Optimizer
AdamW
Learning rate
2e-4
A weight decay
0.01
AdamW (31, 92)
(0.9, 0.99)
AdamW
le-6
Max. grad norm
Batch size
8
Training horizon (generation length)
32
Denoiser (ra) Hyperparameters:
Number of Transformer layers
12
Number of Attention heads
8
Attention head dimension Embedding dimension
64
512
64
Spatial patch size
Reward Termination Model Hyperparameters:
CNN hidden channels
256
CNN output channels
64
LSTM hidden dimension
512
Table 5: Actor-critic hyperparameters.
Description
Value
Optimizer
Adam
Learning rate
2c-4
Adam weight decay
0.01
AdamW (81, 82)
(0.9, 0.99)
Max. grad norm
GAE
GAE y (discount factor)
0.99
Entropy weight
0.001
Imagination actor-critic context length
20
Imagination world model context length Imagination batch size
Imagination generation horizon
32
For Atari, we follow the standard default setting with one exception: sticky actions are replaced by the max-initial no-ops wrapper for randomness, as prior work (Micheli et al., 2023 Cohen et al.) 2025 2024 Allono et al., 2024). Due to computational constraints, we were limited to a single evaluation across the benchmarks. Moreover, our agent's architecture was untested sample-efficiency setting, with prior evidence to support its performance. We therefore adopted a well-established setting with proven track record to mitigate risk. In hindsight, however, given the observed results, this precaution may not have been necessary. Nonetheless, unlike most prior works, we did not use sign rewards or termination on life loss. The Atari hyperparameters are presented in Table 6
17

---

Published as a conference paper at ICLR 2026
Table 6: Atari 100K hyperparameters.
Description
Value
Frame resolution
x
Frame color space (RGB/ grayscale)
RGB
Frame Skip
4
Max random initial no-ops
30
Sticky actions probability
0.0
Terminate on live loss
No
Sign rewards
No
Table Craftium hyperparameters.
Description
Value
Frame resolution
64 x 64
Frame Skip
4
time speed
0 (fixed time of day) True
Sync mode
FPS max. (all excl. ChopTree-vO)
200
FPS max. (ChopTrre-vO)
10
A.3 IMPLEMENTATION DETAILS
Actor Initialization Since exploration difficulty is not central to our evaluation, we initial- ize the actor network's biases with task-specific priors in select Craftium environments to en- courage more effective data collection in the early stages, before actor-critic training begins. In Crafti weaaftiim/eae www.m/Spelleeeee we set the forward action bias to 1 and all others to 0. In we set the dig (chop) action bias to and the rest to O.
Craftinm/ChoppTe
Replay Buffer Sampling Under uniform sampling, training progresses and the replay buffer grows, the probability of sampling recently collected data decreases. To mitigate this, we sample 70% of the batch uniformly, while the remaining 30% are sampled from Beta(3, 1) distribution (sample index is i = [x], Beta(3, 1)).
18

---

Published as a conference paper at ICLR 2026
B RUNTIME ANALYSIS
Table & provides a detailed runtime analysis for the online training of the agent. Table 9 includes imagination throughput details.
Table 8: Breakdown of training stage durations (in milliseconds) for budgets B 16 and B 32 measured on an RTX 4090 GPU.
Stage (Time in ms)
16
B 32
Tokenizer training step
72.5
72.5
World model training step
133.2
133.2
Controller training step
1071
1871
Imagination
673
1326.5
Reward prediction
5.18
5.18
Denoising total (Bx steps)
652.8
1305.6
Single step duration
40.8
40.8
Action computation
6.5
6.5
Denoiser forward
34.3
34.3
Epoch Total Time (sec)
Tokenizer (300 steps)
21.75
21.75
World model (300 steps)
39.96
39.96
Controller imagination (50 steps)
53.55
93.55
Table 9: Imagination throughput for budgets = 16 and B = 32 under the configuration detailed
in Appendix A measured on an RTX 4090 GPU.
B 16
B 32
Frame denoising steps per second
23529.4
23529.4
Frame generations per second 32-frame segments per second
1470.6
735.3
46
23
19

---

Published as a conference paper at ICLR 2026
C ALGORITHM PSEUDOCODE
Algorithm
Agent Training Outline
1:
procedure training.loop
2: repeat
3: collect.experiene trum.steps.to.coleect for representation model update steps do
4:
5:
crain.repreeentttt.co..oooo 6: end for
()
7:
for world model update steps do
8:
()
9:
end for
10:
for controller update steps do Seain.conntraii end for
11:
() (Alg.) )
12:
13:
until stopping criterion met (e.g., target number of epochs)
end procedure
procedure colLect.experienee
16:
Initialize controller state with latest context (real env.)
for
do
18:
Sample action a, <(a (Zce)
Apply the action in the environment: Oc41, + env. env. step(a.) Store experience in replay buffer
20:
21:
if d = then
22:
Reset environment and model context.
23:
end
24: end for
25:
end procedure
26:
procedure creiin.represennttioon model
Sample batch of (independent) observations (frames) o from the replay buffer Compute encoder outputs (o)
28:
Compute decoder outputs (reconstructions) ô Gdee
30:
Compute loss (Eq. 4 Section A.1.1
31:
Update model weights
32:
end procedure
33:
procedure train.
34:
Sample a batch trajectory segments from the replay buffer
35:
Compute latent observation representations 2, (O,) for all observations
36:
Sample corresponding noise sequences z and denoising times 7 Compute noisy sequences z (1-
37:
38:
Compute denoiser outputs ua(z", a) (single forward pass)
Compute denoiser loss (Eq. [
40:
Compute eoward-termination loss (Section A.1.2
Update models weights end procedure
42:
43:
procedure train.controller
44:
Sample a batch trajectory segments from the replay buffer (context)
45:
Initialize controller state and set world model context Generate future trajectory using Horizon Imagination (Alg. I
46:
47:
Compute control objectives (Sections 4 and A.1.3)
48:
Update controller weights
end procedure
20

---

Published as a conference paper at ICLR 2026
D ADDITIONAL RESULTS
Parallel vs. Sequential Generation Quality Atari Figure 7 presents additional generation quality evaluations for the Atari benchmark. Notably, we observe the same overall trends while FVD values are generally lower across configurations due to the visually simpler observation space and dynamics. The MSE values are orders of magnitude smaller, which aligns with the fact that the majority of pixels remain fixed between frames in Atari.
Decay Horizon
(AR)
16
32
ALE/Boxing-/5
ALL/CracyCiimeers
ALL/GophereyS
ALEKingluMesster>
16-
34
12-
33
30-
33
10-
20
10
15
22
54
128
236
18
32
14
125
276
15
22
54
125
54
125
256
Denoising Budget (0)
Denoising Budget
Denoising Budget
kin Budget
ALE/Boxing4S
ALE:CraryClimberu5
ALE/Gogharv55
0.00315
0.0012
0.00010
Eror
0.0226
0.0020
0.0224
0.00050
0.0325
0.00345
0.032
0.0331
0.00040
Denoising
Denoising Budget
Denoising Budget
Denoising Budget
Figure 7: World model generation quality versus denoising steps budget for the subset of Atari games. Each point shows the average FVD/MSE over 512 sampled 33-frame segments, where the first frame was given context and the last 32 were generated conditioned on the recorded actions. dashed vertical line indicates the transition out sub-frame budgets.
Pyramidal Schedule (Diffusion Forcing) Performance Degradation Figure reports results under the same experimental setup Section 5.3 except using the Pyramidal schedule of et al. 2024).. Our findings reveal significant degradation at high budgets, driven by the coupling of budget and decay horizon. Instead of improving with more computation, generation quality collapses the budget grows. Furthermore, unlike our schedule, this approach lacks support for sub-frame budgets.
21

---

Published as a conference paper at ICLR 2026
Decay
(AR)
8
16
32
Craftlum/Splkoovv
CratiumiChogTrse-v0
ALE/Boxinq-VS
ALEXGopher-v5
110
102
340
70
90
70
PVD
300
50-
176
Denaising Budget (8)
Denoising Budget (1)
Denoising Budget (8)
Dinciiing Budget (8)
Crattium/Speloc-v0
CoatummCCaapeeev00
ALE/Boxing-45
ALE/GopheruS
0.0227
0.0500
0.0037
00022
0.0220
Maan Souared
0.0000
0.0224
0.0035
0.004
0.0323
0.0350
0.0034
00012
0.0311
32
128
25c
0.0000
24
128
250
128
236
120
255
Denolsing Budget (B)
Denolsing Budget
Denoising Budgot
Danoising (udget(
Figure 8: World model generation quality versus denoising steps budget when using the Pyramidal schedule Chen et al. (2024). Each point shows the average FVD/MSE over 512 sampled 33-frame segments, where the first frame was given context and the last 32 were generated conditioned on the recorded actions. A dashed vertical line indicates the transition out of sub-frand budgets.
Qualitative Visualizations of Horizon Schedule Configurations To complement the quantita- present qualitative visualizations of the generated sequences used in
tive results in Section 5.3 that analysis, across different configurations of the proposed Horizon schedule. From the large set of possible configurations and samples, we select subset that highlights the most important cases, reflects overall trends, and also provides a visual impression of the environment and the task. We further show ground-truth frames (GT) and their tokenizer reconstructions (Rec) to evaluate repre- sentation quality in isolation from the world model. This enables us to disentangle enrors due to imperfect representations from those caused by dynamics modeling. Representative examples are provided in Figures 9 l6
22

---

Published as a conference paper at ICLR 2026
Rec
p
p-
-
B 82
p= 4
p=1
y-
t=0
t=
t=10
t=15
t=
t=
t= 30
Figure 9: Qualitative visualization of generated sequences in Craftium/CChppTee under different Horizon schedule configurations (B,v), Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with 1 correspond to the autoregressive baseline.
23

---

Published as a conference ionderene paper at ICLR 2026
Rec
B
y
- 32
8 32 v=l
-
2=25 - -
1=0
5
10
15
20
25
Es 30
Figure 10: Qualitative Aieullizaiinn of generated sequences in Graftium um RRom-ve vO under differ- ent an Monizzon schedule configuee configurations v)) Ground-rruth frames (GT) and their vokenizzr recon-
tunuaiionn (Rec) are grovided for ifcrence. Configurattona with v = correspont to the a
gressive Rsselinc.
24

---

Published as a conference paper at ICLR 2026
Rec
p
p-
32
B 1
v= 4
p=1
y=1
t=0
t=
t=10
t=15
t=
t=
t= 30
Figure 11: Qualitative visualization of generated sequences in Craftium wamm/pplelov/ under different Horizon schedule configurations (B,v). Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with 1 correspond to the autoregressive baseline.
25

---

Published as a conference paper at ICLR 2026
Rec
p
p-
32
1
p=4
p=1
y=
t=0
t=
t=10
t=15
t=
t=
t= 30
Figure 12: Qualitative visualization of generated sequences in Craftiumi Sma 1 Room vo under different Horizon schedule configurations (B, Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with 1 correspond to the autoregressive baseline.
26

---

Published as a conference paper at ICLR 2026
Rec
p
144
240
240
240
- 32
B 92 1
4
p
344
240
216
246
y=
t=0
t=
t=10
t=15
t=
t=
t= 30
Figure 13: Qualitative visualization of generated sequences in ALE/Gopher-vé under different Horizon schedule configurations (B, Ground-truth frames (GT) and their tokenizer reconstruc- tions (Rec) are provided for reference. Configurations with = 1 correspond to the autoregressive baseline.
27

---

Published as a conference paper at ICLR 2026
Rec
y-
-
B 32
1
4
p
y=
t=0
t=
t= 10
t=15 t=
t=
t= 30
Figure 14: Qualitative visualization of generated sequences in ALE/ LLE/Boxiigg-ve under different Horizon schedule configurations (B, Ground-truth frames (GT) and their tokenizer reconstruc- tions (Rec) are provided for reference. Configurations with = 1 correspond to the autoregressive baseline.
28

---

Published as a conference paper at ICLR 2026
Rec
p
y
- 32
8 32
1
4
p
y-
1=0
t=5
t=10
t=15
t=
t=
t= 30
Figure 15: Qualitative visualization of generated sequences in ALE/ Ttt//KnnFFMMasteerrae under different Horizon schedule configurations (B, Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with 1 correspond to the autoregressive baseline.
29

---

Published as a conference paper at ICLR 2026
Rec
p.
p-
- 32
B 32
p= 4
p
y
1=0
t=5
t= 10
t=15
t=
t=
t= 30
Figure 16: Qualitative visualization of generated sequences in ALE/CrazyCl wiwryciimmmiis under different Horizon schedule configurations (B,v). Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with 1 correspond to the autoregressive baseline.
30

---

Published as a conference paper at ICLR 2026
E PROOF FOR PROPOSITION 1
E.1 PROOF TO PROPOSITION
Proof. Let p AN Without loss of generality, let p (., so that p(i) which simplifies notations.
Let S N. (p) = Eソ= Pj 0, then by the definition of or it holds that Pr(a(p,-) i) = Pr(w; < a:(p)) = Pr(w < p:)
Otherwise, S,(p) > 0. In that case,
Pr(a(p;))= =) = Pr(we < (p)))) II
Pr(ws 2 az(D))
Pi
(1-ay(p))
Σ k= Pk
Pj
Σ k==
k= Pk
y=1
i-1
Pi
k=n+t 1 Pk
kimi Pk j=1
ス Pk
Σ
Pi
2k Pk
km
kmn Pke
Ekmi
Pk
kmml 1 Pk
k
Pk:
In the last equality, used the fact that Σ Pk 1. The probability of the last action follows from the fact that the probability mass sums to
E.2 PROOF TO PROPOSITIONII
Proof. The first inequality
p all < Pr(A) P Pr(a(p,w) A a(q.w)) Yo E 10,1)
follows from the coupling lemma (Aldous 1983) (Lemma 3.6, p.249) and from Proposition T..
For the second inequality Pr(A) ||a(p) a(g)lli,
u ma) aqq))
min(a(p), (g))) (min(a) (min(ai(p))eo (1))) (i)(mx--(()
Without loss of generality, let p = (1, .N) so that p(i) = i, which simplifies notations. Recall that
A = { app w) # w)q,w)), w ~01001 1) Consider the following partition of [0, 1) N-1
N- U ((LCC)) UU,
L, (w Vj 2: w. 2 uy and w, < 4) C. = [w] d: I My and E 41 w:)
U= (4 1A N -11:w 2 u)
31

---

Published as a conference paper at ICLR 2026
Notice that by the definition of a(., -), it holds that a(p,w) = a(q,w) for all w E [ Ln] uu.
Similarly, Pr(A) Pr(w U Ct). Hence,
N
Pr(A)= Pr(C,)
(u (1) (1 - -u²)
la:(p) a:99) uz)
Σ lo:(p) a:(9))
(p( a(q)ll
32

---

Published as a conference paper at ICLR 2026
F POTENTIAL EXTENSION TO CONTINUOUS ACTION SPACES
Although the stable action sampling mechanism in Horizon Imagination is tailored to discrete action spaces, we propose a direct extension continuous domains that we leave for future investigation. Specifically, consider Gaussian policy with continuous outputs
Pt. oe = T-- aci).
To adapt the method to the continuous case, we consider two modifications. First, the initial sample distribution is drawn from a standard normal, ~ (O, Second, using the reparameterization trick, we define the action mapping as
= -
This construction would minimize action changes when the policy distribution parameters μ,, main fixed, reduce action differences when they vary, and ensure that a(т,w) w)) follows the policy's distribution, in the discrete case. Notably, such extension may also avoid the computational overhead associated with classifier-guidance methods used in prior work.
G DECOUPLED CLEAN AND NOISY ACTOR VARIANT
During additional development experiments on Atari Breakout, we observed that the actor update used in the main paper underperforms this game. We did not observe similar degradation in the environments evaluated in the main experiments. Motivated by this observation, we explored a decoupled actor variant that maintains separate clean and noisy policy copies.
The clean actor is used for interaction with the real environment during data collection, and is up- dated during the imagination phase using the same actor-critic objective in the main method (REINFORCE with advantage estimation and entropy regularization), computed on fully denoised imagined rollouts.
The noisy actor is used exclusively during world-model imagination, where it is invoked at each de- noising step and actions are selected using our stable sampling mechanism. After rollout generation, the noisy actor is trained via cross-entroe cross-entropy so that, given noisy latent observations, it matches the action distribution produced by clean actor on the corresponding clean latents.
Preliminary experiments indicate improved performance on Breakout and promising behavior on a small number of additional Atari games. Due to compute constraints, we did not re-run the full benchmark suite with this variant, and leave comprehensive evaluation future work.
H USE OF LARGE LANGUAGE MODELS (LLMs)
In our work, we used LLMs for the following tasks:
1. Aid and polish writing. This includes revising original drafts of sentences or short para- graphs for improved phrasing.
2. Writing python scripts that produce matplotlib plots from results. This includes ma- nipulating existing results in raw data formats (e.g., ".csv", ".json, images, videos), and quickly drawing high quality plots using matplotlib.
33


Reviews:
[{"review_id": "8wwzvjrDO3", "reviewer": "Reviewer_kptM", "summary": "The paper proposes Horizon Imagination, an on-policy training procedure for reinforcement learning with diffusion world models that denoise multiple future observations in parallel. This is done by introducing a stable discrete-action sampling mechanism to avoid spurious action flips during denoising, and a novel Horizon schedule that decouples the denoising budget from the decay horizon ν, enabling sub-frame budgets and finer control of compute–quality trade-offs. Experiments on Atari 100K and Craftium show the method maintains control performance with only half the denoising budget and improves generation quality under parallel configurations.", "questions": "Do you foresee a principled extension of a(π,ω)to continuous spaces (beyond discretization)?\n\nHow sensitive are returns to policy entropy regularization when using a(π,ω)? Is there a sweet spot where action stability helps most?\n\nFVD/MSE trends favor certain ν/B regimes; can you relate those trends directly to control performance (e.g., correlation analyses across seeds)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed horizon schedule is a neat design: by fixing ν while varying B, it breaks the tight coupling seen in pyramidal schedules, allowing consistent temporal denoising behavior across budgets and enabling sub-frame B < h operation.\n \nThe stable action sampling a(π,ω) for discrete policies is elegant and theoretically justified: action changes between denoising steps are bounded below by total variation distance and above by a derived l_1  term, greatly reducing unnecessary flips during denoising. \n\nSolid empirical analysis: (a) action-consistency experiments demonstrate near-optimal behavior vs. TV lower bound and strong improvements vs. naïve sampling; (b) control performance comparison across ν/B settings; (c) generation quality vs. ν/B via FVD and MSE Clarity.", "weaknesses": "The proposed method, especially the action sampling, is only for discrete action spaces. This limits applicability to many continuous-control tasks \n\nThe paper would benefit if a per-stage runtime analysis and real-time control throughput (fps) comparison were presented to show the improvement. \n\nIt would help to connect the theoretical bound to returns—e.g., does reduced action-flip rate correlate with improved advantage estimates or policy gradient variance?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fUP9Kx9WRE", "reviewer": "Reviewer_irHg", "summary": "This work proposes Horizon Imagination that tackles a common problem in model-based reinforcement learning with diffusion world models. In each step, the policy receives an observation from the world model, which is subsequently used to generate an action. The world model then predicts the next observation. This sequential dependency on the policy leads to a long inference time that is a bottleneck in fast training. Horizon Imagination, therefore, proposes denoising multiple future observations in parallel in conjunction with a stabilization mechanism to account for the changes in the actions and a sampling schedule. The effectiveness of the proposed method is shown on several benchmark environments.", "questions": "- I found it a bit difficult to follow. As a non-expert in Flows/Diffusion models, I was a bit confused with the notation. Is it for every z several denoising time steps sampled in Eq. 1, or is it only one denoising time step per sample?\n\n- Section 4.2 states that Eq. 1 requires knowing all actions; however, the velocity field v_\\theta is expecting actions $a_{<t}$, which are all actions in the history starting from the current time-step $t$. Doesn't this mean that there is actually no need for knowing future actions, as stated in Section 4.2?\n\n- From my understanding, the paper proposes a method that turns the strict autoregressive structure of the world model (i.e., policy gives action, world model predicts observation, ...) into a parallel inference version. However, I don't understand how this autoregressive structure is still retained in the current version? I think the autoregressive prediction structure still needs to hold in general for world models. \n\n- Given that the policy is queried for noisy variables z, aren't there superscripts to the actions in Eq. (1) missing?\n\n- Section 4.3 mentions using the REINFORCE algorithm. It is well-known that the REINFORCE has high variance in the gradient estimates, whereas the reparameterization trick has smaller variance in the gradients. Is there a reason why the reparameterization trick was not used, e.g., in connection with the Gumbel softmax [1] policy representation that allows using the reparameterization trick? \n\n- Is the proposed method also applicable to continuous state-action spaces? What would be necessary in this case? \n\n[1] E. Jang, et al. Categorical Reparameterization with Gumbel-Softmax. ICLR 2017.", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper tackles an interesting and important problem in world models for policy learning.  I think the proposed method is quite important for the RL field.", "weaknesses": "- I think the presentation of the paper can be improved. \n \n - Please also see my questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CfRqlTXK9a", "reviewer": "Reviewer_8WiC", "summary": "To address the low sampling efficiency of existing diffusion world models, especially the quadratic computational cost arising from the product of environment steps and diffusion steps, the authors introduce Horizon Imagination, an algorithm that enables more efficient on-policy sampling of discrete stochastic policies within diffusion world models. The authors make two key improvements over prior diffusion model sampling methods. First, they introduce a horizon sampling schedule that separates the total denoising budget from the decay horizon, allowing flexible and effective configuration. They also propose a novel sampling strategy that stabilizes policy output actions during denoising and prove it preserves the original sampling distribution. Experiments on Atari 100K and Craftium show that Horizon Imagination greatly reduces denoising steps and outperforms previous schedules like pyramidal schedule in generation quality.", "questions": "1. Are there any possible exploratory directions for extending the idea of Horizon Imagination to continuous action spaces?\n\n2. Algorithm 1 only includes the world model sampling process. Could the authors provide a complete pseudocode of the entire training pipeline, including world model training and RL algorithm updates? This would help readers gain a clearer understanding of the overall framework.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The diffusion-based world model studied in this paper is a rapidly growing research area with significant value in both offline data generation and online policy learning. It holds great promise for substantially reducing the cost of real-world interactions.\n\n2. The authors effectively resolve the instability that occurs when diffusion models and policies interact to jointly sample multi-step trajectories by introducing a theoretically grounded stable action sampling method. This approach significantly enhances the quality and stability of long-horizon trajectory generation.\n\n3. Compared with the previous pyramidal schedule, the proposed horizon schedule provides a more general formulation. By decoupling the decay horizon from the denoising budget, it enables more consistent denoising schedules and achieves higher-quality generation.", "weaknesses": "1. The stable action sampling algorithm proposed by the authors is only applicable to discrete action spaces, which limits the applicability of the Horizon Imagination framework in more general environments with continuous action spaces.\n\n2. A key weakness is the limited scope of the experimental comparisons. The control performance results in Section 5.2 are structured as an internal ablation study, comparing the proposed parallel method only against an autoregressive baseline within their own framework. The paper does not benchmark its end-to-end performance against other established world model agents, such as those discussed in the related work. The absence of direct, end-to-end SOTA comparisons makes it difficult to fully assess the proposed method's relative performance and efficiency in the field.\n\n3. Training for 19 or 27 hours on A100 for relatively simple games like Atari seems somewhat costly. Conducting experiments in more challenging environments could make the proposed method more practically valuable.\n\n4. There is an extra quotation mark at the end of line 446.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)