FIRM-WM: State-factorized factual-interventional recurrent modeling for reward-free visual planning

Yilun Wu; Yunjian Zhang; Aobo Li; Mujiangshan Wang; Haitao Wu; Aqiang Zhang

cs.LGSep 19, 2026

FIRM-WM: State-factorized factual-interventional recurrent modeling for reward-free visual planning

Authors: Yilun Wu, Yunjian Zhang, Aobo Li, Mujiangshan Wang, Haitao Wu, Aqiang Zhang

Organizations: Saint Petersburg State University, Saint Petersburg, Russia · Shenzhen Kaihong Digital Industry Development Co., Ltd., Shenzhen, China · Chinese Academy of Sciences (CAS) – Shenzhen Institute of Advanced Technology · Harbin Institute of Technology, Harbin, China ·

Abstract

Reward-free latent world models can learn from offline videos and solve new image--goal tasks by optimizing actions through predicted latent futures. This setting places two demands on the planning state: its coordinates must be comparable with a goal image. Moreover, its dynamics must retain velocity, motion trend, contact, and other history--dependent information beyond those goal coordinates. Offline training creates a second mismatch: each recorded trajectory reveals one factual future, whereas a sampling--based planner compares many actions that were not taken from the same state. We introduce FIRM-WM (Factual--Interventional Recurrent World Model), a compact pixel world model designed around these two gaps. Its recurrent state separates a typed, goal--comparable configuration from a 128-dimensional dynamic fiber used for prediction but excluded from the terminal goal cost. Broad factual trajectories provide state coverage, while common--reset intervention branches provide observed outcomes for alternative action sequences. Before executing each branch, we reset the environment and restore the same recorded values exposed by the environment's state--setting interface. Under matched CEM planning and three independent full-pipeline seeds, FIRM-WM reaches 99.0 $\pm$ 1.0% on TwoRoom, 92.7 $\pm$ 2.1% on Reacher, and 88.0 $\pm$ 3.0% on OGBench-Cube, compared with 89.0%, 88.0%, and 70.0% for LeWM. The deployed model uses 2.98--3.42M parameters and records 2.13--11.60 $\times$ lower planning time on these tasks.

Figures & tables

Method	Pixel deployment	Reward-free WM	Common-reset branches	Planning	Active parameters
RC-aux ( Li et al., 2026 )	✓	✓	–	reachability-aware CEM	$\sim$ 15.7 M †
SAGE ( Cheng et al., 2026 )	✓	✓	–	subgoal prior + CEM	$\sim$ 49.3 M †
LeWM ( Maes et al., 2026 )	✓	✓	–	CEM	$\sim$ 15 M
Fast-LeWM ( Gao and Xu, 2026 )	✓	✓	–	CEM	$\sim$ 18 M
PSG-JEPA ( Yan et al., 2026 )	✓	✓	–	CEM / policy	NR
PhyLatent ( Zeng et al., 2026 )	✓	✓	–	CEM	LeWM backbone

Table 1: Representative pixel world models.

Figure 1: FIRM-WM image-goal planning architecture. The pixel–action history is mapped by the encoder, causal history module, and grounder to the current configuration $c_{t}$ and belief $b_{t}$ . The history-to-fiber initializer constructs $m_{t}$ , giving the recurrent state $s_{t}=(c_{t},m_{t})$ . For each candidate macro-action sequence, the single shared transition recursively updates configuration and fiber for $H$ steps. The shared encoder and grounder map the goal image to $c_{g}$ , and the typed terminal distance $d_{\mathcal{C}}(\hat{c}_{H},c_{g})$ supplies the CEM objective $J(U)$ . The controller executes the first optimized macro action, appends the next observation and executed action to the causal history, and replans. The fiber affects prediction but is excluded from the goal cost.

Method	TwoRoom	Reacher	Push-T	OGBench-Cube
PLDM ( Sobal et al., 2025 )	97.00	78.00	78.00	65.00
DINO-WM ( Zhou et al., 2025 )	100.00	79.00	74.00	86.00
LeWM ( Maes et al., 2026 )	$84.33\pm 4.23$	$82.67\pm 4.42$	$84.67\pm 6.53$	$67.33\pm 5.01$
Sub-JEPA ( Zhao et al., 2026 )	$95.00\pm 2.76$	$84.00\pm 4.00$	$\bm{89.00\pm 5.33}$	$76.33\pm 5.99$
FIRM-WM (ours)	99.0 $\pm$ 1.0	92.7 $\pm$ 2.1	75.7 $\pm$ 4.7	88.0 $\pm$ 3.0

Table 2: Planning success rate (%) across four continuous-control environments. FIRM-WM reports mean

\pm

sample s.d. over three independent full-pipeline training seeds, each evaluated on the same 100 paired trials. All comparison rows are reproduced from Table 1 of Zhao et al. (2026) : its LeWM and Sub-JEPA entries average six evaluation seeds with 50 trajectories per seed, whereas its PLDM and DINO-WM entries are transcribed from Maes et al. (2026) . Bold highlights the leading result(s) in each column; on the near-saturated TwoRoom task, both DINO-WM’s

100.00

and FIRM-WM’s

99.0\pm 1.0

are highlighted.

Figure 2: Control results with optimization variability. Bars show the mean of three independently trained FIRM-WM seeds; error bars are sample standard deviations and circles are individual seeds. LeWM uses its fixed checkpoint on the same paired evaluation trials and matched CEM budget. FIRM improves all three seeds on TwoRoom, Reacher, and Cube; Push-T exposes the persistent-contact boundary.

Environment	FIRM active	FIRM sec./ep.	LeWM sec./ep.	Speedup	Matched CEM budget
TwoRoom	3.41M	$0.416\pm 0.056$	3.224	7.83 $\times$	300 / 30 it. / top-30 / $H{=}5$
Reacher	3.16M	$0.504\pm 0.149$	5.550	11.60 $\times$	300 / 30 it. / top-30 / $H{=}5$
Push-T	2.98M	$1.860\pm 0.243$	3.926	2.13 $\times$	300 / 30 it. / top-30 / $H{=}5$
OGBench-Cube	3.42M	$0.635\pm 0.118$	2.389	3.84 $\times$	300 / 10 it. / top-30 / $H{=}5$
LeWM reported model size	$\sim$ 15M active parameters; time normalized above within each task				identical per row

Table 3: Active deployment size and recorded planner-loop time on the same evaluation host. FIRM seconds are mean

\pm

sample s.d. over three full-pipeline seeds; the LeWM checkpoint is fixed. Speedup averages the three per-seed ratios.

Training condition	Reacher	Cube
Factual only	89.3 $\pm$ 1.5	79.7 $\pm$ 1.5
Intervention only	71.7 $\pm$ 4.6	82.0 $\pm$ 1.7
Factual + intervention (FIRM)	92.7 $\pm$ 2.1	88.0 $\pm$ 3.0
No history fiber †	68.0	–
Mismatched outcomes †	16.0	–
LeWM (fixed paired baseline)	88.0	70.0

Table 4: Full-chain training-source ablations. Values are success (%), mean

\pm

sample s.d. over three independently initialized training runs.

\dagger

marks Reacher-only single-seed diagnostics; LeWM is the fixed paired baseline. All entries use released-original CEM.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Group	Symbol	Meaning
Environment	$t$	Raw environment time index.
	$o_{t}$	Pixel observation at raw time $t$ .
	$a_{t}\in\mathcal{A}$	Raw control applied at time $t$ .
	$h_{t}$	Causal observation–action history available at time $t$ .
	$o^{g}$	Goal image.
Macro rollout	$S$	Number of raw controls in one macro action.

Appendix

Table 5: Notation used throughout the paper. Raw environment time

t

and local macro index

k

are distinct. The symbol

\chi

denotes training-only auxiliary dynamics, while

d_{\mathcal{C}}

denotes the deployment goal cost.

Environment	$\dim(c)$	Goal-comparable configuration	Periodic pairs	$\dim(\chi)$	History frames	$\dim(u_{k})$
TwoRoom	2	agent position	0	2	3	10
Reacher	4	sine/cosine of two joint angles	2	2	3	10
Push-T	6	block orientation and agent/block positions	1	5	3	10
OGBench-Cube	3	fixed affine-normalized block position in 3-D	0	25	3	25

Appendix

Table 6: Typed interfaces used by the shared configuration–fiber transition. History length is measured in observations; every macro action contains five raw environment controls.

Environment	Configuration $c$	$N_{\circ}$	Terminal cost
TwoRoom	$c=p^{\mathrm{a}}=(x^{\mathrm{a}},y^{\mathrm{a}})$	0	$\\|p^{\mathrm{a}}-p_{g}^{\mathrm{a}}\\|_{2}^{2}$
Reacher	$c=(e(\vartheta_{1}),e(\vartheta_{2}))$	2	$\sum_{j=1}^{2}\\|e(\vartheta_{j})-e(\vartheta_{g,j})\\|_{2}^{2}$
Push-T	$c=(e(\vartheta^{\mathrm{b}}),p^{\mathrm{a}},p^{\mathrm{b}})$	1	$\max\{r_{\vartheta},r_{p}\}^{2}$ , where $r_{\vartheta}=\|\operatorname{wrap}(\vartheta^{\mathrm{b}}-\vartheta_{g}^{\mathrm{b}})\|/\tau_{\vartheta}$ and $r_{p}=\\|\Delta p\\|_{2}/\tau_{p}$ (task-normalized)
OGBench-Cube	fixed affine-normalized 3-D block position	0	$\\|c-c_{g}\\|_{2}^{2}$

Appendix

Table 7: Environment-specific configuration and terminal-cost instantiations. Here

p^{\mathrm{a}}

denotes agent position,

p^{\mathrm{b}}

denotes block position, and a subscript

g

denotes the value inferred from the goal image. One periodic scalar is stored as one unit-circle pair

e(\vartheta)=(\sin\vartheta,\cos\vartheta)

; it does not introduce a second physical degree of freedom.

Environment	Candidates	Iterations	Elites	Macro horizon	Raw steps / macro	Goal offset
TwoRoom	300	30	30	5	5	25
Reacher	300	30	30	5	5	25
Push-T	300	30	30	5	5	25
OGBench-Cube	300	10	30	5	5	25

Appendix

Table 8: Released-original-CEM evaluation contract. Candidate action coordinates use the train-split standardizer stored in each checkpoint, search Gaussians are unbounded, and the deployed sequence is the final elite mean. No row changes the solver between FIRM-WM and LeWM.

Environment	Pixel + history frontend	Configuration grounder	Recurrent transition	Active total
TwoRoom	0.374M	1.613M	1.424M	3.411M
Reacher	0.992M	0.828M	1.343M	3.162M
Push-T	0.449M	0.157M	2.371M	2.977M
Cube	0.376M	1.613M	1.434M	3.424M

Appendix

Table 9: Active deployment parameters by module. Pixel + history frontend combines the deployed visual encoder/tokenizer and causal history module. Training-only reconstruction and change-support heads are omitted.

Environment	Anchors	H5 mean	H5 P95	Max	Gate
TwoRoom	128	0	0	0	Pass
Reacher	442	$4.52{\times}10^{-15}$	$1.75{\times}10^{-14}$	$3.55{\times}10^{-14}$	Pass
Push-T	512	0.0613	0.1424	0.9548	Pass
OGBench-Cube	128	0.00352	0.01035	0.06111	Pass

Appendix

Table 10: Common-reset factual-suffix replay on training anchors. Errors are configuration discrepancies in task success-tolerance units after 25 controls (

H{=}5

); “Max” is over all five macro horizons. The fixed gate requires P95

\leq 0.25

at every horizon and Max

\leq 1.0

. This tests restored-interface fidelity, not identity of hidden simulator or solver state.

Module	Initialization	Objective	Updates	LR ( $\times 10^{-4}$ )	Status
TwoRoom
Tokenizer	random	RGB reconstruction + change support	100	$2$	frozen
Grounder	random head	configuration + finite-difference regression	1,000	$3$	frozen
History module	random	causal token-dynamics NLL	500	$3$	frozen
Unified transition	random $\rightarrow$ carried	recursive $\mathrm{F}\rightarrow(\mathrm{F}+\mathrm{I})$ rollout	$3{,}000\rightarrow 3{,}000$	$3\rightarrow 1$	active
Reacher

Appendix

Table 11: Module training contracts. Arrows denote consecutive stages; update counts and learning rates align from left to right. For OGBench-Cube, S1–S3 denote initial module training, joint Tokenizer–Grounder refinement, and height-balanced refinement.

Environment	Active sources by stage	Updates by stage	$N_{\mathrm{rec}}^{(e)}$
TwoRoom	$\mathrm{F}\rightarrow(\mathrm{F}+\mathrm{I})$	$3{,}000\rightarrow 3{,}000$	$6{,}000$
Reacher	$\mathrm{I}\rightarrow\mathrm{I}\rightarrow(\mathrm{F}+\mathrm{I})$	$4{,}000\rightarrow 2{,}000\rightarrow 3{,}000$	$9{,}000$
Push-T	$(\mathrm{F}+\mathrm{I})\times 6\rightarrow\mathrm{SR}$	$5{,}000\rightarrow 3{,}000\rightarrow 3{,}000\rightarrow 5{,}000\rightarrow 10{,}000\rightarrow 3{,}000\rightarrow 2{,}000$	$31{,}000$
OGBench-Cube	$\mathrm{F}+\mathrm{I}$	$3{,}000$	$3{,}000$

Appendix

Table 12: Complete recurrent-transition update schedules. Arrows separate consecutive stages;

\mathrm{F}

and

\mathrm{I}

denote factual and intervention sources. The final stage in every row uses Equation ( 9 ).

Task	Source $\sigma$	$(\kappa,q(k))$	Error multiset $\mathcal{E}(\mathcal{B})$ for $\mathcal{L}_{\mathrm{tail}}$	component weights
TwoRoom	factual $\mathrm{F}$	$(\mathrm{emb},k-1)$	terminal configuration MSE	$(2,2,0.5,0.5,0.5)$
TwoRoom	intervention $\mathrm{I}$	$(\mathrm{emb},k-1)$	terminal configuration MSE	$(2,2,1.0,0.5,0)$
Reacher	factual $\mathrm{F}$	$(\mathrm{rel},0)$	all-horizon maximum joint chord error	$(2,2,0.5,0.5,0.5)$
Reacher	intervention $\mathrm{I}$	$(\mathrm{rel},0)$	all-horizon maximum joint chord error	$(2,2,0.5,0.5,0.5)$
Push-T	factual $\mathrm{F}$	$(\mathrm{emb},k-1)$	terminal configuration MSE	$(2,2,0.5,0.5,0.5)$
Push-T	intervention $\mathrm{I}$	$(\mathrm{emb},k-1)$	terminal configuration MSE	$(2,2,1.0,0.5,0)$

Appendix

Table 13: Exact rollout-loss contract. Each row lists

(w_{\sigma,c},w_{\sigma,\mathrm{disp}},w_{\sigma,\mathrm{tail}},w_{\sigma,\chi},w_{\sigma,\chi 0})

from Equation ( 19 ).

Figure 3: Full-chain training-source ablations. Each panel shows three independently initialized training runs (dots), their mean, and sample s.d. for factual-only, intervention-only, and combined training. The combined pipeline performs best in every run on both tasks. Reacher additionally includes two separate single-seed diagnostics marked

\dagger

. All results use released-original CEM; dashed lines show the task-specific fixed LeWM paired baselines.

Environment	Seed	FIRM	LeWM	$\Delta$ [95% CI]	Rescue / harm	FIRM sec.	LeWM sec.
TwoRoom	1	99	89	$+10$ [5,16]	10 / 0	0.480	3.224
	2	100	89	$+11$ [5,18]	11 / 0	0.378	3.224
	3	98	89	$+9$ [3,16]	10 / 1	0.391	3.224
Reacher	1	95	88	$+7$ [0,15]	11 / 4	0.675	5.550
	2	92	88	$+4$ [ $-4$ ,12]	10 / 6	0.407	5.550
	3	91	88	$+3$ [ $-5$ ,11]	10 / 7	0.429	5.550

Appendix

Table 14: Per-seed paired results.

\Delta

is FIRM minus LeWM in success points; intervals are 95% paired bootstrap intervals. Rescue/harm count trials won only by FIRM / only by LeWM. LeWM timing is fixed within an environment because its checkpoint and paired evaluator are reused.

Environment	Training episodes	Validation episodes	Factual rollout segments	Intervention states	Action branches per state	Branch rollouts	Evaluation trials
TwoRoom	9,000	1,000	16,384	128	32	4,096	100 paired trials
Reacher	9,000	1,000	16,384	512	16	8,192	100 paired trials
OGBench-Cube	9,000	1,000	4,096	128	16	2,048	100 paired trials
Push-T	16,685	2,000	16,384	512	16	8,192	100 paired trials

Appendix

Table 15: Dataset statistics for the final joint recurrent stage. A factual rollout segment contains a three-frame history and its recorded 25-control continuation. Intervention states are the recorded interface states used to initialize action branches; total branch rollouts equal intervention states times action branches per state. All three models in an environment are evaluated on the same 100 paired trials.

World Models for Robotics

World models for robotics predict how scenes and robot states change under actions to support planning, control, and learning.

See all 384 papers

IMWM: Intuition Models Complement World Models for Latent Planning
Jun 1, 2026Baoqi Gao, Ruize Han, Miao Wang +1World Model-Based Planning Goal-Conditioned Planning
Beihang University · Shenzhen University of Advanced Technology
Grounded World Model: Latent Planning with Language Goals
Apr 13, 2026Quanyi Li, Lan Feng, Haonan Zhang +4Language-Conditioned Robot Planning World Models for Robotics
EPFL · NUS
Back to Parsimonious Latents: Learning Task-Centric World Models from Visual Foundations
May 25, 2026Minghao Fu, Fan Feng, Nicklas Hansen +1World Model Learning World Models
University of California, San Diego

Jun 1, 2026J/K move · Enter open · S save

IMWM: Intuition Models Complement World Models for Latent Planning

Baoqi Gao, Ruize Han, Miao Wang, Song Wang

Beihang University · Shenzhen University of Advanced Technology

World Model-Based Planning Goal-Conditioned Planning AI Control Model-Based Planning Latent World Models Latent-Space Planning

Open paper AskPDF arXiv

Planning with a learned latent world model is a promising route to control from raw pixels, but a strong world model alone is not enough. We show this experimentally: even with a perfect world model (operationalized by replacing the learned forward predictor with an idealized rollout of the true environment dynamics), a finite-budget sample-based planner still fails on some tasks, indicating that the bottleneck can lie in search rather than in world-model accuracy. Motivated by this gap, we propose IMWM (Intuition Model + World Model), which pairs the world model with an intuition model trained from demonstrations to recognize promising actions. The two models collaborate through three lightweight components: (i) Retrieval Initialization, which initializes the planner's action proposal from a retrieved demonstration; (ii) Hybrid Cost, which combines the intuition score with the world-model rollout cost; and (iii) a Reliability Gate, which adjusts how much the planner trusts intuition in each setting. Across four pixel-based goal-reaching tasks (Two-Room, Reacher, Push-T, and OGBench-Cube), IMWM has higher mean success than the world-model-only planner on all four, with the largest gains on Two-Room (99.2%, +11.5 percentage points) and OGBench-Cube (94.7%, +28.5 percentage points).

Abstract

World Models for Robotics

Explore similar work

IMWM: Intuition Models Complement World Models for Latent Planning

Grounded World Model: Latent Planning with Language Goals

Back to Parsimonious Latents: Learning Task-Centric World Models from Visual Foundations