DreamFormer: Dream Imitation with a Transformer World Model for Language-Conditioned Robotic Manipulation

Mostafa Kotb; Cornelius Weber; Muhammad Burhan Hafez; Stefan Wermter

cs.ROOct 3, 2026

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

Abstract

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 $o_{t}$ and minimizing the KL divergence between the predicted prior distribution $\hat{z}_{t}$ and the posterior distribution $z_{t}$ . 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\times 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 $(\mathcal{LGM})$ from either a language instruction or the final expert latent state, and the latent plan module $(\mathcal{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 $s_{t}=\{h_{t},\hat{z}_{t}\}$ , with $h_{t}$ 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.

Model-Based Methods
Method	Downsample	Number of Instructions in a Row (1000 chains)
Method	$64{\times}64$	1	2	3	4	5	Avg. Len.
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

\rightarrow

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

ABCD $\rightarrow$ D
Method	Pretraining	Number of Instructions in a Row (1000 chains)
Method	Pretraining	1	2	3	4	5	Avg. Len.
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

\rightarrow

D and ABC

\rightarrow

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

pixel-shift augmentation strength (static / gripper)
Configuration	Number of Instructions in a Row (1000 chains)
Configuration	1	2	3	4	5	Avg. Len.
no shift	78.2% (1.5)	56.0% (1.3)	38.3% (1.6)	25.6% (1.3)	17.0% (1.6)	2.15 (0.07)
5 / 2 pixels	81.1% (1.9)	61.2% (1.8)	45.1% (1.5)	33.4% (1.0)	22.7% (1.8)	2.44 (0.08)
10 / 4 pixels $\dagger$	82.3% (1.7)	64.1% (2.2)	47.1% (2.8)	34.2% (2.0)	24.3% (2.0)	2.52 (0.11)
Component ablations at (5 / 2 pixels)

TABLE III: Ablation study on CALVIN D

\rightarrow

D split. The DreamFormer configuration reported in Tables I and II is marked with

\dagger

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\ldots 7$ ) as context and then autoregressively predicts the remaining 56 latent states ( $t=8\ldots 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$ )	$196\,C$	$\mathcal{O}(196^{2}\,CH)$
CompACT [ 29 ]	DINOv3 + resampler + FSQ	8	4	Fixed ( $C$ )	$8\,C$	$\mathcal{O}(4\cdot 8^{2}\,CH)$
IRIS [ 20 ]	VQ-VAE	16	16	Growing ( $H$ )	$16H$	$\mathcal{O}(16^{2}H^{2})$
DreamFormer (ours)	Categorical VAE	1	1	Growing ( $H$ )	$H$	$\mathcal{O}(H^{2})$

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.

Parameter	Symbol	Value
Batch size	$B$	16
Sequence length	$T$	64
Transformer blocks	–	5
Transformer heads	–	12
Transformer embedding dim.	–	768
Stochastic latent state	$z_{t}$	$64\times 32$ (2048)

TABLE V: World model hyperparameters.

Parameter	Symbol	Value
Batch size	$B$	512
Imagination horizon	$H$	32
Context window	$\mathcal{C}$	16
Latent goal dimensions	$g$	32
Latent plan categories	–	32
Latent plan classes	–	32

TABLE VI: Actor–Critic hyperparameters.

Explore similar work

Dream to Remember: Continual Steering of Generalist Robot Policies with World Models
Oct 9, 2026Bahey Tharwat, Iman Nematollahi, Marvin Alles +3Continual Robot Learning Robotic RL
University of Freiburg · Technical University of Munich
DreamTrue: Action-Faithful Robot World Model with Counterfactual Post-Training
Oct 8, 2026Junyan Li, Ruizhi Li, Yu Liu +6World Models for Robotics Action-Conditioned World Models
NLPR, Institute of Automation, Chinese Academy of Sciences (CASIA) · Amap, Alibaba Group
World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models
Jul 30, 2026Xiangcheng Zhang, Yilun DuWorld Models for Robotics World Model-Based Planning
Harvard University

Oct 9, 2026, cs.ROJ/K move · Enter open · S save

Dream to Remember: Continual Steering of Generalist Robot Policies with World Models

Bahey Tharwat, Iman Nematollahi, Marvin Alles, Xingyuan Zhang+2

University of Freiburg · Technical University of Munich

Continual Robot Learning Robotic RL Model-Based RL Generative Replay

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Dream to Remember: Continual Steering of Generalist Robot Policies with World Models

DreamTrue: Action-Faithful Robot World Model with Counterfactual Post-Training

World Action Planner: Generalizable Decision-Making with Action-Conditioned World Models