Prior diffusion agents imagine frames with a world model over several denoising steps and train a separate policy on these frames. DRIFT uses one Flow Transformer for both, imagining each frame in one denoising step and acting from the same backbone.
TL;DR: One Flow Transformer is both the world model and the agent. The policy learns in imagination on the backbone's own features, trained online from scratch. Each imagined frame takes one denoising step, using 33× fewer FLOPs than prior diffusion agents.
I was happy to see DIAMOND show that diffusion models make excellent world models for online RL. But I think it gets the RL part a bit backward. The diffusion model is only a black-box simulator, and a separate policy network learns from its rollouts. I was heavily interested in one generative model for learning from experience. Dreamer seemed like the right idea, since its policy learns on the world model's own state. But Dreamer is recurrent and works in a mostly discrete latent space.
I wanted a cleaner, modern version of that idea, built on a diffusion transformer. In DRIFT, the RL agent sits directly on top of the Flow Transformer's backbone, the same network that imagines the world, and learns in its continuous latent space.
Making this work was harder than I expected. Reading out the control state from vanilla denoising features gave a bad policy, because those features are weak for control. So I borrowed an idea from self-supervised learning to explicitly shape predictive dynamics, which made the features much better. There were a few more things to take care of, like efficiency and stability, which you can read about in the paper. The first problem later led me to a super cool result in DEM.
Method
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. 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 blue line tracks the imagined rollout. 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.
Craftium, four 3D tasks, same protocol. In SmallRoom the clip is one rollout, not a median of 10, because its short episodes end during the warm-up. The imagined SmallRoom view loses the red target block.
RealImaginedCrafter
Crafter, same protocol.
Results
DreamerV3
DIAMONDpolicy on pixels
DRIFT-symlogpolicy on control state
DRIFTpolicy on control state
Mean
DreamerV31.18
DIAMOND1.46
DRIFT-symlog1.33
DRIFT1.35
11.21.41.6
Median
DreamerV30.44
DIAMOND0.37
DRIFT-symlog0.43
DRIFT0.52
0.30.50.7
IQM
DreamerV30.47
DIAMOND0.64
DRIFT-symlog0.62
DRIFT0.67
0.40.60.8
Human-normalized score on Atari 100k, 26 games, 95% intervals
Policy input: DRIFT and DRIFT-symlog act from the backbone's control state, Horizon Imagination and DIAMOND from pixels.
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, mean return at 100k steps
Task
DIAMONDDIAMOND
Horizon Imag.Horizon Imag.
DRIFT-symlogDRIFT- symlog
DRIFTDRIFT
SmallRoom
−172.6
−40.9
−18.6
−19.4
Room
−546.1
−135.5
−649.6
−123.2
Speleo
−2,047.0
−339.6
669.8
1,727.6
ChopTree
7.3
13.8
31.6
15.3
Craftium, mean return at 100k steps. DRIFT has a higher mean return than DIAMOND and Horizon Imagination on all four tasks, and on the three navigation tasks it passes DIAMOND's final return within 12k steps (learning curves in the paper).
Model-free
PPO4.6
Rainbow4.3
World models without diffusion
DreamerV3 XL14.5
DreamerV210.0
Δ-IRIS9.3
Diffusion and flow world models
DRIFT-symlog10.6policy on control state
Horizon Imag.8.1policy on pixels
DRIFT6.4policy on control state
DIAMOND3.8policy on pixels
051015
Crafter score (%) at 1M steps
Policy input: DRIFT and DRIFT-symlog act from the backbone's control state, Horizon Imagination and DIAMOND from pixels.
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. (a) Next-latent prediction, on a subset of 8 Atari games with an earlier agent setup. (b) 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}
}