Left: prior diffusion agents imagine frames with a world model over several denoising steps and train a separate policy on these frames. Right: DRIFT uses one Flow Transformer for both, imagining each frame in one denoising step and acting from the same backbone.
Abstract
Diffusion and flow models are effective world models for visual reinforcement learning, letting agents learn behaviors in imagination. In current diffusion-based online agents, the world model serves as a black-box simulator. A separate policy network learns from its generated rollouts, so the diffusion model's backbone representations go unused for control. We introduce DRIFT, an agent whose policy learns entirely in imagination on the same Flow Transformer backbone as its world model, trained online from scratch. Three modeling choices make this work. First, denoising alone yields a less effective control state for the actor–critic. We therefore train the backbone to predict the next frame's latent from the control state's summary of the history and the current action, a reconstruction-free objective that makes the control state predictive and improves control. Second, the world model's reward predictions are less accurate on the generated latents than on the real ones. A simple and effective mitigation is to also train the world model on its own rollouts, with targets from the real data. Third, shortcut flow matching reduces imagination time to one denoising step per imagined transition while matching multi-step sampling in rollout accuracy. Imagining in a compact continuous latent space, DRIFT uses 33× fewer FLOPs per imagined transition than prior diffusion-based agents. It outperforms diffusion and flow baselines on Craftium and on Crafter with full-scale rewards, and is competitive on Atari. These results show that one diffusion generative backbone can both imagine and control in online RL.
Method
Left: world-model training. DRIFT denoises the latent of each time step and outputs the control state, from which small heads predict the reward (trophy) and the continuation (flag). A next-latent head also uses the control state and the action to predict the encoder's latent of the next observation. Right: actor–critic learning in imagination. DRIFT generates each next latent in one denoising step, and the actor and critic take the resulting control state as input.
Imagination in one step per frame
Each clip shows the real game on the left and DRIFT imagining the same actions on the right, one denoising step per frame, from 8 real context frames taken mid-episode. For every environment the clip is the median of 10 rollouts by pixel error, so misses stay in. Frames are decoded for display.
Atari 100k, eight games, 32 imagined frames. The imagined side drifts from the real one as small errors add up; a brief punch can be missed in Boxing and the BankHeist maze changes colour early.Crafter, same protocol.
Results
Atari 100k, 26 games, 95% bootstrap intervals over 5 seeds. DRIFT is competitive with DIAMOND (the intervals overlap and DIAMOND has the higher mean); TWISTER, a latent world model without per-seed scores, reports higher mean and median.
Craftium, evaluation return over 100k steps, with the baselines' final returns as dashed lines. DRIFT ends above DIAMOND and Horizon Imagination on all four tasks.Crafter score at 1M steps. DRIFT-symlog, trained on full-scale rewards, is the best diffusion or flow world model, while DreamerV3 XL is higher and sign-reward DRIFT falls below Horizon Imagination.Ablations. Left: next-latent prediction, on a subset of 8 Atari games with an earlier agent setup. Right: training on self-generated rollouts, on all 26 games.
BibTeX
@misc{awal2026drift,
title = {One Flow Transformer for Imagination and Control},
author = {Awal, Rabiul and Jeong, Jinseong and Sikarwar, Ankur and Kordjamshidi, Parisa and Zadaianchuk, Andrii and Rajeswar, Sai and Seo, Paul Hongsuck and Agrawal, Aishwarya},
note = {Under review},
year = {2026}
}
