StarWM: Self-Supervised Trained Attention Routing for Robust World Models

Zeqiang Zhang; Fabian Wurzberger; Maximilian Otte; Daniel Schmid; Sebastian Gottwald; Arne Peter Raulf; Daniel Alexander Braun

cs.CVSep 25, 2026

StarWM: Self-Supervised Trained Attention Routing for Robust World Models

Authors: Zeqiang Zhang, Fabian Wurzberger, Maximilian Otte, Daniel Schmid, Sebastian Gottwald, Arne Peter Raulf, Daniel Alexander Braun

Organizations: Ulm University · German Aerospace Center (DLR)

Abstract

A robust world model must strike the balance between faithfully capturing environmental dynamics and abstracting away from irrelevant content. While reconstruction-based world models ensure faithful supervision, they misallocate representational capacity by pixel area rather than dynamics relevance for visual tasks, which can cause task-irrelevant content to dominate the learned representation. Alternatively, reconstruction-free methods avoid this bias but risk discarding possibly relevant information. We propose StarWM, which uses a cross-attention module trained on self-supervised dynamics to decide where reconstruction applies. A dual-stream decoder then restricts reconstruction to the attended regions, with stop-gradient barriers preventing interference between the two objectives. These components allows reconstruction to supervise the visual content of attended regions without contaminating the latent with non-predictive information. On DeepMind Control with dynamic video backgrounds, default (reward-free) StarWM achieves the strongest performance under random-frame distractors and substantially outperforms reconstruction-based baselines under sequential video. In addition, its reward-augmented variant matches or exceeds reconstruction-free methods on sequential video, achieving the highest overall return across all distractor regimes. Mechanistic probing confirms StarWM preserves state attributes with near-perfect fidelity through long-horizon imagination while systematically discarding distractors.

Figures & tables

Figure 1: StarWM disentangles where-to-attend (X-Attn) from what-to-extract (reconstruction $\hat{o}_{t}$ ). Left: Self-supervised objective signals train a cross-attention module to localize predictable, action-conditional content, while the content of the attended regions is shaped by a reconstruction objective through a dual-stream decoder. Stop-gradient barriers (//) keep the two objectives from interfering. For detailed architecture diagram see Appendix B.1 . Right: Interquartile mean (IQM) aggregated across 6 DMC tasks under sequential (VH) and random (VHR) video distractors, normalized by DreamerV3 clean performance (recovery ratio). StarWM and its variant are the only methods robust to both distractor types.

Figure 2: Evaluation environments. Left: DeepMind Control Suite with three background conditions: clean, Video Hard Sequential, and Video Hard Random. Right: NoisyShape: a geometric shape with controllable attributes (shape, color, size) on DAVIS video backgrounds, used for mechanistic analysis with ground-truth attribute probing; position is randomly re-sampled each step. Arrows denote the four possible actions.

Figure 3: Latent probe results on NoisyShape. Probe accuracy for objective-relevant attributes (color, shape, size) from H0 (encoding) through H15 (15 steps of imagination), and objective-irrelevant attribute (position MSE) at H0 (H15 omitted as ground-truth position at

t{+}15

is independent of the state at

t

; higher MSE indicates correct irrelevant treatment); dashed line = chance level.

Figure 4: Cross-sample composition reveals emergent factorization. (a.) Input frames from sources X and Y, with (b.) each stream’s isolated reconstruction beneath. The residual stream renders only background context, while the dynamics stream renders entity content. (c.) Swapping components across X and Y produces controlled outputs: Y’s entity on X’s background, with position determined by whichever attention mask ( $\bar{A}_{X}$ or $\bar{A}_{Y}$ ) is used. The three components specialize into distinct emergent roles, mapping onto the content and location separation of § 3 without any explicit disentanglement objective.

Figure 5: DMC learning curves. StarWM and StarWM -R are the only methods that perform strongly under both sequential (VH) and random (VHR) video backgrounds.

Figure 6: Learned attention maps on DMC. StarWM ’s cross-attention consistently localizes the agent through 4 queries under both clean and distracting settings.

Figure 7: Ablation study of sources of the attention signal. Self-supervised guidance and gradient isolation are jointly necessary.

Appendix figures & tables41 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: StarWM architecture. Top: comparison with DreamerV3. StarWM splits the encoder output via cross-attention into an attended stream (entering the RSSM) and a residual stream (auxiliary VAE, discarded during imagination). Bottom: detailed architecture. The cross-attention module (red) extracts entity tokens guided by self-supervised losses; the dual-stream decoder (green) composes attended foreground content with residual background content for reconstruction. Blue components are inherited from DreamerV3. Stop-gradient barriers (//) isolate attention parameters from reconstruction gradients.

Component	Hyperparameter	Value
Cross-attention	Number of queries $N$	4
	Number of heads $H$	4
	Head dimension $d/H$	64
	Projected channels $d$	256
	Intercept resolution	$8\times 8$
	Positional encoding	2D sinusoidal

Appendix

Table 1: Hyperparameters of StarWM .

Algorithm	Clean (1M)	VH (3M)	VHR (1M)
StarWM (ours)	35	94	34
Orig DreamerV3	30	104	37
Iso-Dream++	—	103	37
DenoisedMDP	—	37	15
R2-Dreamer	—	26	10

Appendix

Table 2: Average GPU-hours per task on a single NVIDIA H100/A100 (one training run packs 3 random seeds onto the same GPU). Numbers are total wall-clock time to reach the target step budget (1M for Clean and VHR; 3M for VH).

Figure 9: C.1 : Illustration of NoisyShape Environments.

	Small, $r=6.4$	Medium, $r=11.52$	Large, $r=16.64$
Circle	$S=129$ ( $3.14\%$ )	$S=417$ ( $10.18\%$ )	$S=870$ ( $21.24\%$ )
Triangle	$S=53$ ( $1.30\%$ )	$S=172$ ( $4.21\%$ )	$S=360$ ( $8.78\%$ )
Square	$S=82$ ( $2.00\%$ )	$S=265$ ( $6.48\%$ )	$S=554$ ( $13.52\%$ )
Pentagon	$S=97$ ( $2.38\%$ )	$S=316$ ( $7.70\%$ )	$S=658$ ( $16.07\%$ )

Appendix

Table 3: Parameters for shapes.

	H0 (Encoding)				H15 (Imagination)
Model	Color	Shape	Size	Position MSE $\uparrow$	Color	Shape	Size
DreamerV3	99.92 $\pm$ 0.01	69.37 $\pm$ 2.41	98.92 $\pm$ 0.93	0.006 $\pm$ 0.001	99.73 $\pm$ 0.11	28.95 $\pm$ 1.65	87.49 $\pm$ 14.10
+ dual-stream only	100.0 $\pm$ 0.00	40.01 $\pm$ 1.13	51.55 $\pm$ 2.49	0.068 $\pm$ 0.001	99.98 $\pm$ 0.00	25.04 $\pm$ 0.10	33.31 $\pm$ 0.29
+ SS only	99.97 $\pm$ 0.01	74.15 $\pm$ 2.59	99.65 $\pm$ 0.12	0.006 $\pm$ 0.000	99.90 $\pm$ 0.02	29.04 $\pm$ 1.08	97.69 $\pm$ 0.92
+ dual-stream + SS	100.0 $\pm$ 0.00	38.37 $\pm$ 1.30	50.20 $\pm$ 2.75	0.069 $\pm$ 0.000	99.99 $\pm$ 0.01	24.86 $\pm$ 0.16	33.28 $\pm$ 0.29
+ attention only	99.97 $\pm$ 0.01	92.95 $\pm$ 1.21	99.92 $\pm$ 0.02	0.061 $\pm$ 0.001	99.94 $\pm$ 0.01	56.22 $\pm$ 4.35	99.71 $\pm$ 0.13

Appendix

Table 4: D.1 : Latent probe accuracy on Noisy Shape. H0 = encoding step; H15 = 15 steps of open-loop imagination. Each ablation isolates one component of StarWM .

		Walker Walk	Walker Run	Cheetah Run	Cartpole	Finger Spin	Ball-in-cup
Clean	DreamerV3 (Original Paper)	966	797	846	843	886	975
	DreamerV3 (Pytorch)	954.40 $\pm$ 3.15	749.93 $\pm$ 10.64	984.54 $\pm$ 3.08	840.10 $\pm$ 31.89	554.74 $\pm$ 97.56	964.65 $\pm$ 1.53
	StarWM	944.06 $\pm$ 2.57	686.81 $\pm$ 39.45	762.77 $\pm$ 54.33	863.66 $\pm$ 5.62	559.35 $\pm$ 92.24	948.07 $\pm$ 9.50
VH	DreamerV3	327.39 $\pm$ 68.25	95.36 $\pm$ 40.38	73.30 $\pm$ 8.46	160.13 $\pm$ 12.17	217.50 $\pm$ 119.63	118.26 $\pm$ 23.78
	DenoisedMDP	93.14 $\pm$ 24.69	39.70 $\pm$ 2.91	60.93 $\pm$ 25.55	161.14 $\pm$ 7.17	53.96 $\pm$ 36.57	168.34 $\pm$ 46.43
	Iso-Dream++	98.16 $\pm$ 12.49	43.79 $\pm$ 6.89	28.01 $\pm$ 14.61	139.92 $\pm$ 5.77	6.35 $\pm$ 0.71	15.76 $\pm$ 4.05

Appendix

Table 5: D.2.1 : DMC evaluation returns. VH: sequential video (3M steps). VHR: random video (1M steps). Best per (task, setting) in bold .

Panel of Figure 10 — Figure 10: D.2.2 : Aggregate statistical analysis on DMC under video distractors, following Agarwal et al. (2021) . Each row reports VH, VHR, and the combined VH $+$ VHR pool, computed over six DMC tasks and three random seeds (18 runs per method per setting; 36 in the pool).

Figure 11: E.1.1 : Default StarWM ’s attention drifts on Cartpole VH; reward prediction recovers it. Each row shows the input observation followed by the attention maps for the four queries. Default StarWM (top right) concentrates attention near the track endpoints without query specialization. StarWM -R (bottom right) localizes the cart and pole accurately, with queries specializing to distinct object regions.

Symbol	Configuration	SS	RSSM grad.	Recon. grad.
★	StarWM (full)	✓	—	—
✩1	no detach A/B, with SS	✓	✓	✓
✩2	no detach A/B, no SS	—	✓	✓
✩3	only detach B, no SS	—	✓	—
✩4	only detach B, with SS	✓	✓	—

Appendix

Table 6: E.2 : Configurations evaluated in the gradient isolation analysis. Each row indicates which gradient pathways reach the attention parameters. detach A blocks RSSM gradients (via

\bar{A}_{t}

in entity tokens); detach B blocks reconstruction gradients (via

\bar{A}_{t}

in the dual-stream decoder); SS denotes self-supervised attention guidance (inverse dynamics + contrastive). A check mark indicates the gradient or signal reaches the attention parameters.

Figure 12: E.2 : Probe accuracy on NoisyShape under five gradient isolation configurations. NoisyShape probe accuracy is invariant to gradient isolation choice.

		Walker Walk		Cheetah Run
Symbol	Configuration	Clean	VHR	Clean	VHR
★	StarWM (full)	952.37 $\pm$ 3.06	872.94 $\pm$ 50.99	752.37 $\pm$ 34.84	466.03 $\pm$ 53.00
✩1	no detach A/B, with SS	810.66 $\pm$ 157.94	181.08 $\pm$ 17.72	641.66 $\pm$ 52.40	181.75 $\pm$ 31.91
✩3	only detach B, no SS	943.42 $\pm$ 18.18	135.02 $\pm$ 20.74	880.96 $\pm$ 6.10	199.37 $\pm$ 94.74

Appendix

Table 7: E.2 : Gradient isolation ablation on DMC. Episode return (mean over 3 seeds) for Walker Walk and Cheetah Run under Clean and VHR settings. Configurations are defined in Table 6 . Best per column in bold.

Figure 13: E.2 : Learning curves for the gradient isolation ablation on DMC. Episode return over training for the three configurations of Table 7 . Removing detach B causes attention to be co-opted by reconstruction gradients; removing SS leaves RSSM gradients as the only attention signal, insufficient to disambiguate foreground from background under video distractors. Full StarWM combines detach B with SS guidance.

Figure 14: E.2 : Gradient isolation determines attention focus on DMC. Attention maps from the four learned queries on Walker Walk under VHR, comparing the three configurations of Table 7 . Top row: input observation. Rows 2–5: attention map for each of the four queries. Left (✩1: SS+RSSM+Recon., no detach): attention collapses to a nearly uniform distribution. Middle (✩3: RSSM only, detach B without SS): attention shows weak focus (1/4 queries specializing to agent body in VH). Right (★ StarWM : SS with detach A/B): attention is sharply localized on the agent’s body parts (3/4 queries specializing to agent body in VH).

		Clean		VHR
Symbol	Configuration	Entropy $\downarrow$	Max $\uparrow$	Entropy $\downarrow$	Max $\uparrow$
★	StarWM (full)	0.14	0.92	0.15	0.94
✩1	no detach A/B, with SS	3.90	0.04	4.15	0.02
✩3	only detach B, no SS	3.12	0.14	1.27	0.52

Appendix

Table 8: E.2 : Attention map statistics under three gradient isolation configurations on Cheetah Run. Average entropy and maximum attention weight of query

Q_{0}

, computed over evaluation trajectories. Lower entropy and higher max weight indicate sharper, more focused attention. The maximum possible entropy is

\log(64)\approx 4.16

(uniform over the 64 spatial tokens), and the uniform floor for max attention is

1/64\approx 0.016

Perturbation	Return @ 1M
$\lambda_{\mathrm{inv}}=0.3$ ( $3\times$ reduction)	925 $\pm$ 1
$\beta_{\mathrm{bg}}$ : annealing removed entirely	918 $\pm$ 12
$\lambda_{\mathrm{inv}}=3.0$ ( $3\times$ increase)	856 $\pm$ 34
$\beta_{\mathrm{bg}}^{\mathrm{final}}=2.0$ ( $4\times$ increase)	787 $\pm$ 77

Appendix

Table 9: E.3 : One-at-a-time hyperparameter perturbations ( StarWM , Walker Walk VHR, matched 1M steps, mean

\pm

SE, 3 seeds).

Figure 15: E.3.1 : Inverse dynamics is the primary attention signal. Probe accuracy on the RSSM latent for full StarWM (inverse + contrastive), inverse-only, and contrastive-only variants on NoisyShape. Inverse dynamics alone is sufficient for near-perfect attribute retention; contrastive alone fails to localize the foreground.

Figure 16: E.3.1 : Episode return for full StarWM (inverse + contrastive) versus inverse-only across Clean, VH (sequential video), and VHR (random video). Contrastive supervision is essential under VHR, neutral under Clean and under VH.

$N$	Walker Walk	Cheetah Run	Cartpole Swingup
1	170.48 $\pm$ 24.59	181.22 $\pm$ 63.81	191.54 $\pm$ 6.46
2	928.73 $\pm$ 34.39	461.71 $\pm$ 15.09	268.62 $\pm$ 60.34
4	902.31 $\pm$ 37.20	452.77 $\pm$ 22.76	308.00 $\pm$ 12.04
8	585.14 $\pm$ 72.84	387.97 $\pm$ 78.67	394.90 $\pm$ 98.07

Appendix

Table 10: E.3 : Effect of query count

N

on DMC VH (episode return at 3M steps, mean

\pm

SE, 3 seeds).

Figure 17: E.3 : Return curves across varying query counts

N

on DMC VH (mean

\pm

SE, 3 seeds).

Figure 18: E.4 : Latent space probing under predictable object motion. Left: Position MSE across a 15-step imagination horizon. When position follows a deterministic trajectory, StarWM successfully tracks and rolls out spatial coordinates well below chance level. Right: Background identification accuracy at the encoding step (H0). DreamerV3 suffers from severe objective-irrelevant contamination, while StarWM perfectly drops background identity.

Stage	Model	Color	Shape	Size	Position MSE	BG ID $\downarrow$
H0 (Encoding)	DreamerV3	100.00 $\pm$ 0.00	99.70 $\pm$ 0.03	99.96 $\pm$ 0.04	0.0003 $\pm$ 0.0002	82.44 $\pm$ 2.07
H0 (Encoding)	StarWM	100.00 $\pm$ 0.00	99.88 $\pm$ 0.06	99.99 $\pm$ 0.01	0.0008 $\pm$ 0.0002	5.02 $\pm$ 0.17
H15 (Imagination)	DreamerV3	99.98 $\pm$ 0.03	97.46 $\pm$ 1.56	99.99 $\pm$ 0.00	0.0034 $\pm$ 0.0002	—
H15 (Imagination)	StarWM	100.00 $\pm$ 0.00	99.71 $\pm$ 0.18	99.99 $\pm$ 0.01	0.0150 $\pm$ 0.0046	—
Chance level		25.00	25.00	33.33	0.0711	5.00

Appendix

Table 11: E.4 : Probe accuracy under predictable position. Position follows uniform linear motion with boundary reflection. Background ID is a 20-way classification over DAVIS clips. H0 = encoding; H15 = 15 steps of open-loop imagination.

Environment	Model	Color	Shape	Size	Pos. MSE $\uparrow$
H0 (Encoding)
Shape (fixed)	DreamerV3	100.00	100.00	100.00	—
Shape (clean)	DreamerV3	100.00 $\pm$ 0.01	99.87 $\pm$ 0.05	99.96 $\pm$ 0.01	0.002
Shape (clean)	StarWM	100.00 $\pm$ 0.00	100.00 $\pm$ 0.00	100.00 $\pm$ 0.00	0.072
NoisyShape	DreamerV3	99.92 $\pm$ 0.01	69.37 $\pm$ 2.41	98.92 $\pm$ 0.93	0.006
NoisyShape	StarWM	99.99 $\pm$ 0.00	99.92 $\pm$ 0.02	100.00 $\pm$ 0.00	0.066

Appendix

Table 12: E.5 : Probe accuracy across NoisyShape variants. H0 = encoding step; H15 = 15-step open-loop imagination. All values: mean

\pm

std over 3 seeds. Accuracies in %, position in MSE.

Figure 19: E.5 : Latent space probing across NoisyShape variants. DreamerV3’s performance degrades as irrelevant complexity increases, while StarWM is invariant to distraction complexity.

Panel of Figure 20 — Figure 20: F.1 : Small-object stress test. Left: Learning curves for Cheetah Run and Walker Walk with the agent rescaled to half its default linear size (one quarter pixel area), under Clean, VH (3M steps), and VHR (1M steps) settings. Right: Illustration of the environment setup.

Figure 21: F.2 : Latent probe results on NoisyTwoShape. Probe accuracy for color, shape of the large and small object across H0 to H15 imagination horizon, and position MSE at H0. Higher MSE on position indicates correct irrelevant treatment (positions are re-sampled each step). Dashed lines indicate per-object chance levels for position MSE, computed via Monte Carlo simulation over each object’s position distribution under its size-dependent margin: 0.038 (large) and 0.150 (small).

Figure 22: F.2 : StarWM attention maps on NoisyTwoShape. Columns show the input image followed by the four query attention maps.

Figure 23: F.3 : Foreground texture ladder. The same underlying state rendered across levels for the four texture IDs (replacing color). Only foreground appearance changes: L1 uses external procedural textures, L2 uses crops from DAVIS background videos, and L3 alpha-blends these crops into the background ( $\alpha{=}0.5$ , feathered edges).

Figure 24: F.3 : Latent probe accuracy across the ladder (mean

\pm

SE, 3 seeds, 400k steps). Solid lines/filled markers denote the encoding step (H0); dashed lines/open markers denote 15-step open-loop rollouts (H15). Grey dashed lines show chance performance. For position MSE (right panel), higher values indicate that task-irrelevant position details are successfully filtered out. StarWM stays at chance, whereas DreamerV3 encodes exact position.

		texture-id		shape		size		pos. MSE
Level	Model	H0	H15	H0	H15	H0	H15	H0
L1	StarWM	1.000	1.000	0.984	0.866	1.000	1.000	0.065
L1	DreamerV3	0.999	0.999	0.598	0.259	0.996	0.987	0.011
L2	StarWM	0.999	0.999	0.896	0.579	0.999	0.999	0.067
L2	DreamerV3	0.997	0.993	0.600	0.262	0.996	0.980	0.010
L3	StarWM	0.959	0.929	0.354	0.251	0.786	0.698	0.065

Appendix

Table 13: F.3 : Probe accuracy on the texture ladder (3 seeds, 400k steps). H0 = encoding step, H15 = 15-step rollout. Chance is 0.25 for texture and shape, 0.33 for size; chance position MSE is

\approx

0.071.

Setting	StarWM	DreamerV3	R2-Dreamer
Walker VH (3M steps)	893 $\pm$ 41	799 $\pm$ 53	895 $\pm$ 41
Walker VHR (1M steps)	890 $\pm$ 17	394 $\pm$ 34	118 $\pm$ 13

Appendix

Table 14: F.3 : Low-contrast Walker Walk return (mean

\pm

SE, 3 seeds). VH = coherent video background, VHR = per-step randomized frames.

Figure 25: F.3 : Return curves for low-contrast Walker Walk (mean

\pm

SE, 3 seeds).

Region	Frame share	Attention share	Concentration	Shuffled null
Gripper tip ( $r=5$ px)	1.9%	16.6%	8.69 $\times$	2.83 $\times$
Active block	0.7%	4.2%	6.09 $\times$	2.56 $\times$
Static blocks	1.4%	2.1%	1.52 $\times$	2.03 $\times$
Gripper + all blocks	3.3%	18.8%	5.63 $\times$	2.48 $\times$

Appendix

Table 15: F.3 : Spatial attention concentration on CALVIN. Measured on frames where a block moves (

n=204

). Active blocks are identified by ground-truth movement.

Figure 26: F.3 : Attention routing on CALVIN. Input frame and corresponding attention maps from four queries. Attention targets the controllable objects and gripper over background structures.

Figure 27: F.4 : PassiveShapeEnv . Five consecutive steps under each background setting. The active object $A$ changes attributes based on actions and moves at each step; the passive object $B$ (cyan circle) maintains fixed attributes throughout the episode. Top: per-step randomized background frames. Bottom: sequential video background.

	$b$ _colour (acc)		$b$ _ $x$ ( $R^{2}$ )		$b$ _ $y$ ( $R^{2}$ )
Background	H0	H15	H0	H15	H0	H15
random-frame	1.000 $\pm$ .000	1.000 $\pm$ .000	0.80 $\pm$ .09	0.78 $\pm$ .09	0.84 $\pm$ .04	0.76 $\pm$ .08
sequential video	0.998 $\pm$ .000	0.994 $\pm$ .003	0.13 $\pm$ .03	$-0.23$ $\pm$ .02	0.08 $\pm$ .02	$-0.18$ $\pm$ .04

Appendix

Table 16: F.4 : Latent probe recovery for passive object

B

(mean

\pm

SE, 3 seeds, 212k steps).

b

_colour is 4-way classification accuracy (chance = 0.25);

b

_position reports held-out

R^{2}

Figure 28: F.4 : Probe results for passive object

B

(mean

\pm

SE, 3 seeds, 212k steps). Object identity reaches ceiling accuracy across both regimes. Position is recovered on randomized backgrounds, but drops to or below

R^{2}{=}0

under sequential video.

Setting	StarWM	StarWM -R	DreamerV3	R2-Dreamer
Reacher VH (3M steps)	391 $\pm$ 136	789 $\pm$ 33	585 $\pm$ 146	947 $\pm$ 3
Reacher VHR (1M steps)	690 $\pm$ 70	657 $\pm$ 1	450 $\pm$ 37	66 $\pm$ 9

Appendix

Table 17: F.4 : DMC Reacher episode returns (mean

\pm

SE, 3 seeds). VH = temporally coherent video background, VHR = per-step randomized frames.

Figure 29: F.4 : DMC Reacher performance. Left: observation frames showing the two-link arm and static target disc against DAVIS backgrounds (VH: coherent video, VHR: per-step random frames). Right: learning curves (mean $\pm$ SE, 3 seeds). On VH, reward augmentation ( StarWM -R) recovers default StarWM ’s deficit. On VHR, StarWM and StarWM -R lead while R2-Dreamer drops.

Task-relevant factor	Episode-static & distinctive?	Competing signal?	Default StarWM outcome
Passive $B$ position, random bg	✓	None	Encoded ( $R^{2}\!\approx\!0.8$ )
Passive $B$ position, seq. bg	✓	Video	Suppressed ( $R^{2}\!\leq\!0.13$ )
Reacher target, VHR	✓	None	Solved (690)
Reacher target, VH	✓	Video	Low return (391); benefits from reward
Ball in flight (Ball-in-Cup)	✗	None	Requires reward routing

Appendix

Table 18: F.4 : Determinants of task-relevant object retention. Default routing retains features that remain static within an episode and distinct across episodes, provided no competing signal (such as coherent video) already satisfies that criterion.

$k$	Supervision window	Attn. mass	$\times$ Uniform	Shape	Size
0	Reference (no delay)	0.846	11.2 $\times$	0.999	1.000
1	Adjacent (mismatched)	0.439 $\pm$ .033	6.6 $\times$	0.385 $\pm$ .027	0.476 $\pm$ .029
1	Shifted (matched)	0.824 $\pm$ .026	12.3 $\times$	0.998 $\pm$ .001	1.000 $\pm$ .000
2	Adjacent	0.352 $\pm$ .093	5.3 $\times$	0.654 $\pm$ .153	0.722 $\pm$ .143
2	Shifted	0.755 $\pm$ .010	11.3 $\times$	0.997 $\pm$ .000	0.999 $\pm$ .000
4	Adjacent	0.357 $\pm$ .009	6.2 $\times$	0.448 $\pm$ .039	0.504 $\pm$ .028

Appendix

Table 19: F.5 : Delayed action effects at 180k steps (mean

\pm

SE, 3 seeds). Object coverage is

\approx

6.7% of the frame;

\times

uniform denotes attention concentration. Color decodability remains near ceiling (

\approx

1.00) across all settings and is omitted. The

k{=}0

row serves as an un-delayed reference point probed at its final checkpoint (

\approx

145k steps).

Explore similar work

TaskSense: Focusing on What Matters in World Models
Aug 6, 2026SM Mazharul Islam, Manfred HuberWorld Model Learning World Models
1The University of Texas at Arlington
Contrastive World Models
Aug 27, 2026Bonnie LiWorld Model Learning Model-Based RL
Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling
Mar 4, 2026Tal Daniel, Carl Qi, Dan Haramati +5World Model Learning Object-Centric World Models
Carnegie Mellon University · UT Austin · Brown University +2

Aug 6, 2026J/K move · Enter open · S save

TaskSense: Focusing on What Matters in World Models

SM Mazharul Islam, Manfred Huber

1The University of Texas at Arlington

World Model Learning World Models Visual RL

Open paper AskPDF arXiv

World models for visual control typically learn compact latent states by reconstructing observations, implicitly encouraging representations to preserve information across the entire visual input. However, task-relevant content often occupies only a small fraction of the observation, while background clutter and distractors consume valuable representational capacity. This mismatch between visual reconstruction and control objectives biases latent representations to model task-irrelevant visual content, diluting learning signals for control-relevant features and severely degrading downstream performance under visual distractions. We introduce TaskSense, a task-centric world modeling framework that enforces task relevance before latent encoding through a differentiable stochastic spatial attention mechanism conditioned on the previous latent state. To steer attention toward control-relevant regions, we augment training with an auxiliary inverse-dynamics objective. Rather than reconstructing the full observation, the world model reconstructs only the attended regions, encouraging latent representations to preserve task-relevant information while discarding irrelevant visual content. The decoder is further conditioned on the sampled attention map, enabling consistent reconstruction despite stochastic attention. Compared with the DreamerV3 baseline, TaskSense maintains competitive performance on the DeepMind Control Suite while consistently outperforming DreamerV3 on the Distracting Control Suite, demonstrating substantially improved robustness to visual distractions. Qualitative analysis further confirms that the learned attention, guided by inverse-dynamics supervision, consistently localizes control-relevant regions while suppressing irrelevant visual content.

Abstract

Explore similar work

TaskSense: Focusing on What Matters in World Models

Contrastive World Models

Latent Particle World Models: Self-supervised Object-centric Stochastic Dynamics Modeling