Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models

Lior Cohen; Ofir Nabati; Kaixin Wang; Navdeep Kumar; Shie Mannor

cs.LGFeb 8, 2026

Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models

Authors: Lior Cohen, Ofir Nabati, Kaixin Wang, Navdeep Kumar, Shie Mannor

Organizations: Technion · Microsoft Research

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.

Figures & tables

Figure 1: Example of generation instabilities observed in Craftium/ChopTree-v0 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.

Figure 2: Comparison of the Pyramidal schedule ( Chen et al., 2024 ) and the proposed Horizon schedule (transposed). Horizon fixes the decay horizon ( $\nu=3$ ), yielding consistent schedules across budgets, whereas in the Pyramidal schedule the decay horizon drifts with budget, as the two are entangled, leading to degraded generation quality at higher budgets.

Panel of Figure 3 — Figure 3: Empirical study of the average number of action changes under various settings.

Panel of Figure 4 — Figure 4: 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$ ). A dashed horizontal line denotes Atari human-level performance.

Panel of Figure 5 — Figure 5: Actor–critic performance comparison between the proposed stable action sampling method and the naive baseline. Each baseline is evaluated over 5 seeds. Curves show the mean and standard deviation, smoothed by a moving average (window size $15$ ). A dashed horizontal line denotes Atari human-level performance.

Panel of Figure 6 — Figure 6: World model generation quality versus denoising steps budget. Each point shows the average FVD/MSE over 512 sampled 33-frame segments, where the first frame was given as context and the last 32 were generated conditioned on the recorded actions. A dashed vertical line indicates the transition out of sub-frame budgets.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Module	Output Shape
Encoder
Input	$3\times 64\times 64$
Conv2d(3, 1, 1)	$64\times 64\times 64$
ResnetBlock	$128\times 64\times 64$
ResnetBlock	$128\times 64\times 64$
Conv2d(3, 2, 0)	$128\times 32\times 32$

Appendix

Table 1: The encoder and decoder architectures of the representation model. “Conv2d(a,b,c)” represents a 2d-convolutional layer with kernel size

a\times a

, stride of

b

and padding

c

. “GN” represents a GroupNorm operator with

32

groups,

\epsilon=1e-6

and learnable per-channel affine parameters. “Upsample” represents a repeat-interleave operation that doubles the spatial dimensions, followed by a Conv2d(3,1,1) layer.

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	25

Appendix

Table 2: Agent hyperparameters.

Description	Value
Optimizer	AdamW
Learning rate	2e-4
AdamW weight decay	0.05
AdamW ( $\beta_{1},\beta_{2})$	(0.9, 0.95)
Batch size	32
Max. grad norm	1

Appendix

Table 3: Representation model (tokenizer) hyperparameters.

Description	Value
Optimizer	AdamW
Learning rate	2e-4
AdamW weight decay	0.01
AdamW ( $\beta_{1},\beta_{2})$	(0.9, 0.99)
AdamW $\epsilon$	1e-6
Max. grad norm	1

Appendix

Table 4: World model hyperparameters.

Description	Value
Optimizer	AdamW
Learning rate	2e-4
AdamW weight decay	0.01
AdamW ( $\beta_{1},\beta_{2})$	(0.9, 0.99)
Max. grad norm	1
GAE $\lambda$	0.95

Appendix

Table 5: Actor-critic hyperparameters.

Description	Value
Frame resolution	$64\times 64$
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

Appendix

Table 6: Atari 100K hyperparameters.

Description	Value
Frame resolution	$64\times 64$
Frame Skip	4
time speed	0 (fixed time of day)
Sync mode	True
FPS max. (all excl. ChopTree-v0)	200
FPS max. (ChopTree-v0)	10

Appendix

Table 7: Craftium hyperparameters.

Stage (Time in ms)	$\mathbf{B=16}$	$\mathbf{B=32}$
Tokenizer training step	72.5	72.5
World model training step	133.2	133.2
Controller training step	1071	1871
$\llcorner$ Imagination	673	1326.5
$\llcorner$ Reward prediction	5.18	5.18
$\llcorner$ Denoising total ( $B\times$ steps)	652.8	1305.6

Appendix

Table 8: Breakdown of training stage durations (in milliseconds) for budgets

B=16

and

B=32

measured on an RTX 4090 GPU.

	$\mathbf{B=16}$	$\mathbf{B=32}$
Frame denoising steps per second	23529.4	23529.4
Frame generations per second	1470.6	735.3
32-frame segments per second	46	23

Appendix

Table 9: Imagination throughput for budgets

B=16

and

B=32

under the configuration detailed in Appendix A , measured on an RTX 4090 GPU.

Panel of Figure 7 — 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 as context and the last 32 were generated conditioned on the recorded actions. A dashed vertical line indicates the transition out of sub-frame budgets.

Panel of Figure 8 — Figure 8: World model generation quality versus denoising steps budget when using the Pyramidal schedule of Chen et al. (2024) . Each point shows the average FVD/MSE over 512 sampled 33-frame segments, where the first frame was given as context and the last 32 were generated conditioned on the recorded actions. A dashed vertical line indicates the transition out of sub-frame budgets.

Figure 9: Qualitative visualization of generated sequences in Craftium/ChopTree-v0 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 10: Qualitative visualization of generated sequences in Craftium/Room-v0 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 11: Qualitative visualization of generated sequences in Craftium/Speleo-v0 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 12: Qualitative visualization of generated sequences in Craftium/SmallRoom-v0 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 13: Qualitative visualization of generated sequences in ALE/Gopher-v5 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 14: Qualitative visualization of generated sequences in ALE/Boxing-v5 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 15: Qualitative visualization of generated sequences in ALE/KungFuMaster-v5 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Figure 16: Qualitative visualization of generated sequences in ALE/CrazyClimber-v5 under different Horizon schedule configurations $(B,\nu)$ . Ground-truth frames (GT) and their tokenizer reconstructions (Rec) are provided for reference. Configurations with $\nu=1$ correspond to the autoregressive baseline.

Reinforcement Learning

Reinforcement learning studies how agents learn sequential decisions from rewards and interaction, balancing exploration, performance, safety, and adaptation.

See all 1,742 papers

DriftWorld: Fast World Modeling through Drifting
Jul 16, 2026Susie Lu, Haonan Chen, Weirui Ye +1World Models for Robotics World Model-Based Planning
Massachusetts Institute of Technology · Harvard University
Persistent Robot World Models: Stabilizing Multi-Step Rollouts via Reinforcement Learning
Mar 26, 2026Jai Bardhan, Patrik Drozdik, Josef Sivic +1World Models for Robotics World Models
Czech Institute of Informatics, Robotics and Cybernetics, Czech Technical University in Prague
JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning
May 13, 2026Jing Yu Lim, Rushi Shah, Zarif Ikram +4World Model Learning Joint-Embedding Predictive Architecture
National University of Singapore · usc

Jul 16, 2026J/K move · Enter open · S save

DriftWorld: Fast World Modeling through Drifting

Susie Lu, Haonan Chen, Weirui Ye, Yilun Du

Massachusetts Institute of Technology · Harvard University

World Models for Robotics World Model-Based Planning Action-Conditioned World Models Robotic Policy Evaluation Efficient Inference for World Action Models Model-Based Planning+1

Open paper AskPDF arXiv

Predictive world models enable robots to plan by imagining the outcomes of their actions, but their value for control hinges on generating many rollouts quickly. This creates a bottleneck for diffusion-based world models: multistep sampling makes each rollout expensive, limiting large-scale action search at inference time. We introduce DriftWorld, an action-conditioned world model based on drifting generative models. Rather than denoising iteratively at inference, DriftWorld learns an action-conditioned drift during training, allowing it to generate future frames from the current observation and a candidate action sequence in a single forward pass at 30+ fps, which is 17x faster on average than diffusion based baselines. We evaluate DriftWorld on standard vision-based robotic manipulation benchmarks, including Bridge-V2, RT-1, Language Table, Push-T, and Robomimic. By producing rollouts that are both accurate and fast, DriftWorld achieves state-of-the-art decision-making performance with far less inference time than diffusion-based world model baselines. Beyond online control, DriftWorld can also serve as an offline simulator for ranking real-world robot policies, with rollout-based scores correlating with ground truth at up to 0.99. These results show that drifting models are a strong fit for robot world modeling, where fast, high-quality imagination directly supports planning and policy evaluation.

Abstract

Reinforcement Learning

Explore similar work

DriftWorld: Fast World Modeling through Drifting

Persistent Robot World Models: Stabilizing Multi-Step Rollouts via Reinforcement Learning

JEDI: Joint Embedding Diffusion World Model for Online Model-Based Reinforcement Learning