🌅 Horizon Imagination (HI)

May 20, 2026 · View on GitHub

Efficient On-Policy Rollout in Diffusion World Models
📄 OpenReview · 🚀 Quickstart · 📝 Citation

Horizon Imagination improves the imagination efficiency of diffusion-based world models, enabling effective generations with sub-step budgets and efficient on-policy generation for training lightweight controllers in imagination. HI proposes a sampling schedule that improves over the popular Pyramidal schedule (Diffusion Forcing).

This repository contains the training and evaluation code used in the ICLR 2026 paper.

🌿 Branches

main — stable, up-to-date code reflecting the latest tested improvements
dev — active development and integration of new features
feat/* — individual feature branches, merged into dev when complete
iclr-2026 — frozen snapshot of the code used to reproduce the results in the paper

🚀 Quickstart

Requirements

🐳 Docker
🖥️ NVIDIA GPU with recent drivers

This project works best with Docker, as it automatically installs complex dependencies such as Craftium on a tested system setup. Running without Docker is possible by following the dockerfile setup steps (Linux only) but requires manual installation of additional dependencies.

🐳 Docker Setup

Build the image

docker build -t horizon-imagination .

Run the container

docker run --gpus all -v .:/horizon_imagination -it horizon-imagination /bin/bash

Make sure to run the above command from the project root directory.

🎮 Training an Online Agent

Training uses a client–server setup.
Run tmux (within the Docker container), and execute the following command to start the server:

1️⃣ Start the environment server

portal-env start <env_name> -b mm

Supported environments:

ale — Atari (ALE)
craftium — Craftium (Minecraft-like)

Example:

portal-env start ale -b mm

2️⃣ Launch training (new tmux tab)

Start a new tmux tab (Ctrl-b c) and launch a training run via

python experiments/train_agent_online.py \
  -b <env_name> \
  -g <game_id>

Example:

python experiments/train_agent_online.py \
  -b ale \
  -g ALE/Boxing-v5

🧱 Project Structure

horizon_imagination/
├── horizon_imagination/   # Core models, imagination logic, training code
├── experiments/           # Training and evaluation entry points
├── config/                # Experiment configurations
├── scripts/               # Utilities (setup, analysis, plotting)
├── outputs/               # Local artifacts (logs, checkpoints)
├── results/               # Official results reported in the paper
└── README.md

🧰 Common Issues

🧑‍🔧 Docker file ownership

Docker may change the ownership of the generated output files on some systems. Fix with:

sudo chown -R 1000:1000 .

🧾 Citation

If you use this code, please cite:

@inproceedings{
  cohen2026horizon,
  title={Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models},
  author={Lior Cohen and Ofir Nabati and Kaixin Wang and Navdeep Kumar and Shie Mannor},
  booktitle={The Fourteenth International Conference on Learning Representations},
  year={2026},
  url={https://openreview.net/forum?id=Obefq4k8iG}
}