Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models
Fuente:
arXiv
Enregistré dans:
| Auteurs principaux: | , , , , |
|---|---|
| Format: | Preprint |
| Publié: |
2026
|
| Sujets: | |
| Accès en ligne: | |
| Tags: |
Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
|
| _version_ | 1866912910334230528 |
|---|---|
| author | Cohen, Lior Nabati, Ofir Wang, Kaixin Kumar, Navdeep Mannor, Shie |
| author_facet | Cohen, Lior Nabati, Ofir Wang, Kaixin Kumar, Navdeep Mannor, Shie |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2602_08032 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models Cohen, Lior Nabati, Ofir Wang, Kaixin Kumar, Navdeep Mannor, Shie Machine Learning 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. |
| title | Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2602.08032 |