Grid or Nodes? Probing Object Memory in World Models Under Partial Observability

Justin McMillen; Yasin Yilmaz

cs.ROOct 13, 2026

Grid or Nodes? Probing Object Memory in World Models Under Partial Observability

Authors: Justin McMillen, Yasin Yilmaz

Organizations: University of South Florida Tampa, FL, USA

Abstract

A robot's world model must remember objects that leave the camera's view. FloWM, a learned world model, keeps them in a spatial grid of features that it moves with five built-in velocities. A node memory instead keeps each object's position and velocity as numbers. We compare them in FloWM's own block world, with an oracle-fed Kalman tracker as the node memory. How the grid is read changes its score: a single-cell probe finds 66% of out-of-view blocks, while FloWM's own probe 89%. FloWM also scores 8 points higher in view although no frame is shown, which we attribute to its write-back of its own predictions. Out of view, the tracker reaches FloWM's range in 14KB against 5.6MB, with no probe. On diagonal motion, outside its five velocities, FloWM finds 11% of blocks even after fine-tuning, while the unchanged tracker keeps 93%. Faster blocks defeat both until one more allowed speed takes the tracker to 94%, where FloWM's built-in shifts would need 25 copies of its grid. A learned node memory must still replace the oracles and handle velocity under noise, long gaps, and the bounce.

Figures & tables

Panel of Figure 1 — Figure 1: (a) The room from above. Solid arrows show motion in the standard setting, and dashed arrows the diagonal and speed-2 (double head) motion that the varied setting adds. (b) One prediction step of FloWM’s grid memory. (c) The tracker, a node memory.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Part	Tracker	FloWM
Finding blocks in a frame	Oracle: true masks	Learned: write transformer
Block position in a frame	Oracle: true position plus one-pixel noise	Learned: write transformer
Block appearance	Oracle: mean color inside the true mask	Learned: 256 features per cell
Agent’s own motion	Hand-set: from actions	Built in: shift and rotation by actions
Motion between steps	Hand-set: Kalman predict step	Built in: five velocity copies
Allowed speeds	Fitted: snap each axis to 0 or 1	Built in: the five velocities

Appendix

Table 2: Where each part of the two memories comes from. Oracle parts are given by the simulator. Hand-set parts are written by us, with no training. Fitted parts are estimated from noisy oracle positions on training videos of the standard setting, with thresholds chosen on the tuning videos. Built in parts are fixed by FloWM’s design. Learned parts are trained with FloWM’s frame-prediction loss, except the probe, which is trained to read positions. The table describes the tracker with allowed speeds 0 and 1.

Panel of Figure 3 — Figure 3: The protocol for one block. Both memories receive the same context phase. FloWM receives the frames, and the tracker receives oracle detections from the same frames. At every predicted step, each memory names one grid cell for the block, and the answer is correct only if it is the block’s exact cell. A sighting is a context frame in which the block covers at least 12 pixels. We split answers by whether the block lies inside the camera’s view at the scored step and by the steps since the block’s last sighting.

All blocks
Tracker variant	Out of view
Exact positions, no speed snapping	93.7
One-pixel noise, no speed snapping	47.1
One-pixel noise, speed snapping (the tracker)	90.8
Exact positions, speed snapping	94.3
Blocks sighted at least twice $\dagger$

Appendix

Table 3: Out-of-view accuracy (%) of tracker variants on the standard test videos. The row marked “the tracker” is the tracker of the main text. The upper rows score all 85,602 out-of-view answers. The lower rows score only the 83,202 out-of-view answers for blocks sighted in at least two context frames. A restart predicts a block forward from the tracker’s estimate at its last sighting, with no further matching.

	Probe	Standard	Varied videos				MSE
Model	trained on	videos	All	Built-in	Diagonal	Speed 2	Standard	Varied
Released	Standard	86.9	16.3	75.8	3.6	4.3	0.000649	0.0148
Released	Varied ∗	–	6.2	22.3	2.2	3.1
Fine-tuned on varied, 20k steps	Varied ∗	84.5	20.1	46.8	13.4	15.0	0.00135	0.00747
From scratch on standard, 150k	Standard ∗	84.8	–	–	–	–	0.000824	0.0158
From scratch on varied, 150k	Varied ∗	30.5	13.1	20.7	10.5	11.8	0.00592	0.00790

Appendix

Table 4: FloWM on the varied setting, as overall accuracy (%) over answers in view and out of view. The standard-video column is on 1,024 fresh standard videos, which include the 256 test videos, and every score in it comes from heads trained on the standard setting while the model generated rollouts. The varied-video columns are on the 1,024 varied test videos and use the heads named in the second column. MSE is on 1,000 clips with 70 predicted frames. Built-in means one cell per step along one axis. Heads marked ∗ were trained while the model generated rollouts, and the other heads were trained on saved memory states. Probe scores carry no interval.

Positions	Allowed speeds	All	Built-in	Diagonal	Speed 2	speed 2	IDF1
						Diagonal and
One-pixel noise	0, 1 and 2	93.4	91.0	94.1	93.9	93.9	93.6
One-pixel noise	0 and 1 (the tracker)	33.3	89.9	93.1	4.9	21.7	89.4
One-pixel noise	no speed snapping	56.5	47.1	54.2	59.5	58.5	88.3
Exact	no speed snapping	97.0	95.1	97.1	97.5	97.4	94.2
Exact	0, 1 and 2	97.0	95.4	97.2	97.3	97.3	95.6

Appendix

Table 5: The tracker on the 1,024 varied test videos: out-of-view accuracy (%) by motion type, and IDF1 (%). Each configuration ran once, after the code and settings were frozen and the run was registered.

	Probe				Diagonal and
Model	trained on	Built-in	Diagonal	Speed 2	speed 2
Released	Standard	72.7	3.5	4.2	4.1
Released	Varied	18.6	2.1	3.1	2.9
Fine-tuned on varied, 20k steps	Varied	39.1	11.0	13.0	12.6
From scratch on varied, 150k	Varied	16.8	8.2	9.4	9.1

Appendix

Table 6: FloWM on the 1,024 varied test videos: out-of-view accuracy (%) by motion type, on the same answers as Table 5 . The first row comes from per-answer outputs. The other rows come from accuracies logged per motion type, and their last column assumes the first row’s answer counts.

Comparison	Difference	95% interval
Standard test videos, out of view: the tracker minus FloWM read by the convolutional probe	$+$ 2.1	0.5 to 3.7
FloWM, in view minus out of view, matched on steps since the last sighting	$+$ 7.5	6.7 to 8.4
The tracker, in view minus out of view, same matching	$-$ 0.1	$-$ 1.2 to 0.9
FloWM, a block after it enters view during prediction minus before, test videos	$+$ 4.6	3.7 to 5.5
The tracker minus FloWM, out of view, 10 to 19 steps since the last sighting	$+$ 6	4 to 8
The tracker minus FloWM, out of view, 100 to 139 steps since the last sighting	$-$ 10	$-$ 16 to $-$ 4

Appendix

Table 7: The main-text comparisons that have an interval, in points of accuracy. Each difference is the first quantity minus the second, so a negative value means the first is lower. Intervals are 95% bootstrap intervals.

World Model Learning

World model learning builds predictive models of how environments change under actions to support planning, simulation, and decision-making.

See all 414 papers

Scope-WM: Scoped Computation for Efficient Visual World Models
Sep 27, 2026Chunzheng Li, Zesheng Jia, Hongda Zhang +4World Models for Robotics World Models
Soochow University · Fudan University
Better Slots, Better Worlds: Representation Quality & Robustness in Object-Centric World Models
Aug 12, 2026Shukrullo Nazirjonov, Sai Prasanna, Anna Manasyan +1World Models World Model-Based Planning
University of Tuebingen · Max Planck Institute for Intelligent Systems
MemoWM: How World Models Change What Agents Need to Remember
Oct 7, 2026Bingfan Zeng, Zhisheng Chen, Chenbo Sang +6Memory Compression Agent Memory Management
South China University of Technology · University of the Chinese Academy of Sciences · Beijing University of Posts and Telecommunications +5

Sep 27, 2026J/K move · Enter open · S save

Scope-WM: Scoped Computation for Efficient Visual World Models

Chunzheng Li, Zesheng Jia, Hongda Zhang, Jiaying Tang+3

Soochow University · Fudan University

World Models for Robotics World Models World Model Learning World Model-Based Planning Model-Based Planning

Open paper AskPDF arXiv

Visual world models enable robotic planning by predicting future observations, but dense latent-state propagation and sample-intensive trajectory optimization incur high inference latency and peak memory usage, limiting real-time deployment on resource-constrained platforms. Existing sparse world-model acceleration methods either rely on unguided token sparsification, which may discard planning-relevant information and restrict achievable sparsity, or introduce heavy auxiliary modules and cumbersome multi-stage training pipelines. In this work, we present Scope-WM, an efficient visual world model that scopes computation to prediction-relevant latent regions and promising action sequences. Scope-WM distills prediction relevance into a lightweight action-conditioned selector and applies full dynamics prediction only to a compact subset of selected tokens. It updates the remaining tokens using a compact summary of foreground states and their changes, allowing the background to perceive foreground dynamics without costly token-to-token interactions. During planning, Scope-WM preserves and reuses high-quality action sequences discovered during the initial MPC search, focusing subsequent search under reduced rollout budgets. The resulting pipeline requires only a one-off selector distillation followed by a single joint training stage for the sparse world model. On the challenging Push-T task, Scope-WM reduces peak GPU memory usage and planning time to

18.1\%

and

14.3\%

of those of dense DINO-WM, respectively, corresponding to a

6.97\times

planning speedup, while maintaining competitive task performance. Further evaluations across five diverse visual planning tasks demonstrate the general applicability of Scope-WM. Code is available at https://github.com/ChunZheng2022/Scope-WM.

Abstract

World Model Learning

Explore similar work

Scope-WM: Scoped Computation for Efficient Visual World Models

Better Slots, Better Worlds: Representation Quality & Robustness in Object-Centric World Models

MemoWM: How World Models Change What Agents Need to Remember