Diffusion models learned to draw the future in exquisite detail. The trouble, for anyone using them to train a reinforcement learning agent, is that they draw it one painstaking denoising pass at a time. A new paper from researchers at the Technion and Microsoft Research proposes a way to sketch many frames of that future at once, and reports it can cut the compute of on-policy rollouts roughly in half without giving up control performance.
The work, by Lior Cohen, Ofir Nabati, Navdeep Kumar and Shie Mannor of the Technion, together with Kaixin Wang of Microsoft Research, targets a bottleneck that has quietly limited how useful high-fidelity diffusion world models can be for the thing world models were supposed to be good at: training controllers cheaply.
The problem
World models let an agent learn from imagined experience rather than expensive interaction with the real environment. Diffusion-based world models are attractive because they produce the highest-fidelity image and video predictions available. But that fidelity is expensive. Every predicted observation is generated through a multi-step denoising process, and classic "imagination" interleaves policy decisions with world model predictions one step at a time. The result, as the authors put it, is inherently sequential: generate a frame, ask the policy what to do, generate the next frame, and so on. The paper notes that prior diffusion world model agents reported significant overhead, with one detailed runtime analysis showing sequential imagination dominating the total cost. The field's response has largely been to scale models up, which only makes each denoising step slower, running against what control needs: a lightweight controller capable of fast, low-power inference.
The method
Horizon Imagination, or HI, attacks the sequential burden directly by denoising multiple future observations in parallel rather than finishing one before starting the next. It is designed for discrete stochastic policies, and the authors describe it as training-agnostic: it applies to any pre-trained world model that conditions on time at the observation level.
Two ingredients make the parallel scheme work. The first is a stabilization mechanism for action sampling. When many frames are denoised at once, the policy must produce actions at every denoising step, from pure noise to clean samples. Naively drawing a fresh action at each step, especially from a high-entropy policy, causes the chosen action to flip around constantly even when the policy itself is fixed, corrupting the rollout. HI introduces a "stable" sampling method that reduces those unnecessary changes; the authors prove that if the action distribution does not change between steps, the sampled action stays put.
The second ingredient is a new sampling schedule the authors call the Horizon schedule. It builds on the pyramidal schedule of prior diffusion-forcing work, which denoises near-future frames before far-future ones. The key move is disentangling two things earlier schedules tangled together: the total denoising budget (how many steps you spend) and the decay horizon (how quickly denoising progresses across the sequence). In the pyramidal schedule those are coupled, so pushing the budget higher drifts the effective horizon and degrades generation. The Horizon schedule fixes the decay horizon independently, keeping schedules consistent across budgets and, critically, allowing arbitrary integer budgets, including sub-frame budgets where the total number of denoising steps is smaller than the number of frames generated. Standard autoregressive generation cannot do this; it requires at least one step per frame.
The results
The authors validate HI on a subset of Atari 100K and on Craftium, a 3D benchmark, each evaluated over five seeds. Their headline claim is that the agent maintains the control performance of the autoregressive baseline while using only half the denoising budget. Concretely, with a decay horizon of four, a sub-frame budget of 16 denoising steps sustains full performance on rollouts that generate 32 frames, half the steps the autoregressive baseline requires. They further report that a single denoising step per observation generally suffices; only Craftium's ChopTree, the most visually complex environment tested, benefits from a larger budget.
The efficiency shows up in the runtime tables. Measured on an RTX 4090, imagination throughput at budget 16 reaches 46 thirty-two-frame segments per second, versus 23 at budget 32, a clean doubling. A controller imagination epoch of 50 steps drops from 93.55 seconds at budget 32 to 53.55 seconds at budget 16. The sub-frame regime is not a marginal tweak; it roughly halves the dominant cost.
On generation quality, the authors sweep decay horizon and budget from fully autoregressive to highly parallel regimes, scoring 512 sampled 33-frame segments with FVD and MSE. Parallel generation proves consistently advantageous. They also report a telling contrast: under the older pyramidal schedule, generation quality collapses as the budget grows, because budget and decay horizon remain coupled. The Horizon schedule instead holds steady. One nuance worth flagging is that higher budgets can improve perceptual quality (FVD) while MSE rises, suggesting the generated sequences drift further from ground truth even as they look better.
Why it matters
The significance is less about a new state-of-the-art score and more about changing the cost curve of a promising but expensive paradigm. Diffusion world models have the best generative fidelity going, but their sequential denoising has made on-policy imagination a compute sink. By decoupling budget from horizon and stabilizing action sampling, HI turns denoising steps into a dial that can be turned down below one-per-frame while keeping the controller intact. That the method is training-agnostic and works with any observation-time-conditioned pre-trained world model makes it a plausible drop-in rather than a bespoke architecture.
What to watch
The evaluations cover a subset of Atari 100K and Craftium, not the full suite, and the policies are discrete, so continuous control and larger, more chaotic environments remain open questions. The FVD-versus-MSE divergence is worth tracking, since a model that looks right but drifts numerically could mislead a controller in ways that surface only in long-horizon tasks. With code released, the near-term signal will be whether other diffusion world model agents adopt sub-frame budgets as a default, and whether the halved-cost result holds up as environments and model sizes grow.