DreamFormer: Dream Imitation with a Transformer World Model for Language-Conditioned Robotic Manipulation
Authors: Mostafa Kotb, Cornelius Weber, Muhammad Burhan Hafez, Stefan Wermter
Organizations: Knowledge Technology Group, Department of Informatics, University of Hamburg, 20146 Hamburg, Germany · Mathematics Department, Faculty of Science, Aswan University, 81528 Aswan, Egypt · School of Electronics and Computer Science, University of Southampton, Southampton SO17 1BJ, UK
We introduce DreamFormer, a model-based agent that acquires language-conditioned, multi-task skills by imitating expert demonstrations within the latent imagination of a learned world model. DreamFormer first learns a task-agnostic Transformer world model from unstructured play data, then acquires task-specific behaviors by optimizing an intrinsic reward that aligns agent-generated rollouts with expert demonstrations in latent space. Since the policy is trained on-policy inside imagination, it is exposed to its own errors during training, mitigating the covariate shift inherent to offline behavioral cloning. To make long-horizon imagination affordable, DreamFormer encodes a high-resolution multi-view robotic observation into a single input token, avoiding both spatial downsampling and the multi-token representations used by prior Transformer world models. On the long-horizon CALVIN benchmark, DreamFormer outperforms LUMOS, the comparable model-based agent, on single-environment evaluation (2.52 vs 2.34 average tasks completed per chain of five) while keeping imagination rollouts tractable. Against HULC, the behavior cloning baseline, it nearly doubles performance on zero-shot transfer to an unseen environment (1.30 vs 0.67), indicating that dynamics learned by the world model transfer more readily than a directly cloned policy. This is consistent with the role attributed to internal models in biological agents, where a model of environment dynamics supports behavior in situations not previously encountered.
Figures & tables
Fig. 1: Transformer world model is trained by reconstructing the observation ot and minimizing the KL divergence between the predicted prior distribution z^t and the posterior distribution zt . To accommodate high-resolution, multi-view observations, a trainable ResNet-18 is employed as the visual encoder, and the capacity of the categorical latent space is doubled to 64 categories, each with 32 classes (64×32) .
Fig. 2: DreamFormer learns goal-conditioned behaviors within the latent space of a learned world model. (a) During training, imagined trajectories are unrolled from an expert state, initialized with a key–value cache over a context window of preceding steps. Goals are encoded by the latent goal module (LGM) from either a language instruction or the final expert latent state, and the latent plan module (LPM) supplies a high-level behavioral prior: the plan is drawn from the plan recognizer, which sees the entire expert trajectory, while the plan proposal is trained to match it from the current state and goal alone. The actor is trained by maximizing an intrinsic reward that rewards alignment between imagined and expert latent trajectories. (b) During inference, the actor acts based on the current latent state st={ht,z^t} , with ht computed over a sliding context window of recent real observations, guided by a language command. The latent plan comes from the plan proposal rather than the recognizer, conditioning the actor on the intended behavior.
Fig. 3: CALVIN Benchmark . (Left) the four environments A, B, C, D. (Right) long-horizon, language-conditioned evaluation with 5 tasks in a row.
Method
Downsample
Number of Instructions in a Row (1000 chains)
64×64
1
2
3
4
5
Avg. Len.
Model-Based Methods
LUMOS [ 10 ]
✓
80.7% (1.0)
59.3% (1.4)
42.6% (1.8)
30.7% (1.2)
21.1% (0.8)
2.34 (0.05)
DreamFormer (ours)
✗
82.3% (1.7)
64.1% (2.2)
47.1% (2.8)
34.2% (2.0)
24.3% (2.0)
2.52 (0.11)
Model-Free Methods
GCBC [ 17 ]
✗
64.7% (4.0)
28.4% (6.2)
12.2% (4.1)
4.9% (2.0)
1.3% (0.9)
1.11 (0.3)
TABLE I: Performance of all methods on CALVIN D → D split. Methods are grouped by model class (model-based and model-free). Bold marks the best result within each group. Parentheses give the standard deviation over three seeds
Method
Pretraining
Number of Instructions in a Row (1000 chains)
1
2
3
4
5
Avg. Len.
ABCD → D
MCIL [ 23 ]
✗
37.3%
2.7%
0.2%
0.0%
0.0%
0.40
HULC [ 25 ]
✗
88.9%
73.3%
58.7%
47.5%
38.3%
3.06 (0.07)
DreamFormer (ours)
✗
89.1%
74.0%
60.1%
49.1%
38.4%
3.11 (0.05)
RT-1 [ 42 ]
ImageNet
84.4%
61.7%
43.8%
32.3%
22.7%
2.45
TABLE II: Performance of all methods on CALVIN ABCD → D and ABC → D splits. Methods are grouped by pretraining regime (w/o pretraining and w/ pretraining). Bold marks the best result within each group. Parentheses give the standard deviation over three seeds; the remaining baselines report single-run values
TABLE III: Ablation study on CALVIN D → D split. The DreamFormer configuration reported in Tables I and II is marked with † .
Fig. 4: Open-loop rollouts of the world model. A trajectory from a held-out validation set. The world model observes the first 8 frames ( t=0…7 ) as context and then autoregressively predicts the remaining 56 latent states ( t=8…63 ), conditioned only on the ground-truth action sequence. Rows show the ground truth (GT) and decoded predictions (Model) for the static and gripper views; the vertical rule marks the end of the context window, and frames are subsampled every eight steps. The arm pose and scene layout remain accurate across the full horizon.
Method
Tokenizer
Design choices
Rollout cost ( H steps)
Tokens
Predict
KV cache
Max cached
Attention
/frame
steps/frame
type
tokens
cost
DINO-WM [ 32 ]
ViT patches (DINOv2)
196
1
Fixed ( C )
196C
O(1962CH)
CompACT [ 29 ]
DINOv3 + resampler + FSQ
8
4
Fixed ( C )
8C
O(4⋅82CH)
IRIS [ 20 ]
VQ-VAE
16
16
Growing ( H )
16H
O(162H2)
DreamFormer (ours)
Categorical VAE
1
1
Growing ( H )
H
O(H2)
TABLE IV: Sequence-model memory and attention cost for an imagination rollout of horizon H , derived from the configuration reported in each paper. Predict steps/frame is the number of sequence-model passes needed to produce one imagined frame; Max cached tokens is the size of the key–value cache, which determines memory; Attention cost counts query–key interactions over the rollout, assuming key–value caching and matched model width and depth. C is the length of the fixed context window (3 frames for DINO-WM, 4 for CompACT). CompACT values are for its navigation configuration.
Generalist robot policies are trained with imitation learning and are bounded by the quality and coverage of their demonstrations. Reinforcement learning can improve and extend them, either by fine-tuning the policy weights or by training a small module to steer the input noise of a frozen generative policy. When tasks arrive sequentially, both approaches suffer from catastrophic forgetting. Fine-tuning on a task degrades prior tasks, and steering only shifts forgetting to the shared steering module. We propose Dream to Remember (D2R), which continually steers a frozen generalist policy toward new tasks, while retaining its earlier skills. As the steering module learns a new task, a frozen pretrained world model rehearses prior tasks through imagined rollouts, anchoring the module to its earlier behavior without a replay buffer or new data. When the world model is faithful enough, D2R learns the new task entirely in imagination. Otherwise, it learns the task in the environment while using the world model only to rehearse prior tasks. On CALVIN, D2R improves eight tasks in sequence entirely in imagination, raising average success from 63.7% to 74.8%, while LoRA fine-tuning on the same sequence drops it to 11.7%. With a frozen off-the-shelf generalist policy and world model, π0 and A2World, D2R learns unseen LIBERO-90 tasks in the environment while keeping π0's performance on the LIBERO suites it was fine-tuned on, even though the world model is not accurate enough for policy optimization. Finally, on a real Franka robot, D2R improves three tasks in sequence entirely in imagination while retaining the earlier improvements.
Bahey Tharwat, Iman Nematollahi, Marvin Alles +3
University of Freiburg · Technical University of Munich
We present DreamTrue, a multi-view, cross-embodiment robot world model for action-faithful and physically plausible video prediction. Training such a model on existing robot datasets faces two obstacles: imprecise calibration can impair action following, while limited coverage of unsuccessful interactions can bias predictions toward successful outcomes. To improve action following across embodiments, we render action trajectories into image-space conditions and introduce offline geometric calibration to align these conditions with the target videos. To broaden interaction coverage, we introduce counterfactual post-training, modifying recorded action trajectories and generating future videos under a wider range of actions and contact configurations. To provide feedback on these predictions without paired ground-truth futures, we construct a human-annotated video dataset covering robot, object, and interaction defects and use it to train an embodied video reward model. Its scores guide reinforcement-learning post-training toward more physically plausible interaction outcomes. On AgiBot, DreamTrue attains state-of-the-art action following, while reducing the human-assessed interaction defect rate from from 48.12% to 6.25%. Notably, our model ranks first in the world model track of the AgiBot World Challenge 2026. The project page can be found at https://brave-eai.github.io/DreamTrue.
Junyan Li, Ruizhi Li, Yu Liu +6
NLPR, Institute of Automation, Chinese Academy of Sciences (CASIA) · Amap, Alibaba Group
Building generalizable agents for diverse applications remains a fundamental challenge. While imitation learning-based policies succeed in specific training environments, they often fail to generalize to novel scenes and tasks. In this work, we propose World Action Planner, a robot planning system that leverages the reasoning capabilities of Vision-Language Models (VLMs) and the physical grounding of a multi-task pose-image conditioned world model. Our system enables an agent to propose initial action plans and iteratively refine them via optimization and search, reasoning over imagined world model rollouts. We demonstrate that our approach achieves superior performance across compositional tasks, new layouts, and zero-shot generalization scenarios, significantly outperforming state-of-the-art end-to-end policy models such as VLAs and WAMs. Project website at worldactionplanner.github.io