Rollout-Decoded Reconstruction for Long-Horizon Prediction in Latent World Models

Rishi Shah; Rishav Shrestha

cs.LGAug 25, 2026

Rollout-Decoded Reconstruction for Long-Horizon Prediction in Latent World Models

Authors: Rishi Shah, Rishav Shrestha

Organizations: E3A Healthcare

Abstract

A latent world model trains its decoder on latents anchored to observations, then deploys it on the model's own free-running rollout, hundreds of steps past the last observation. Rollout-Decoded Reconstruction (RDR) closes this gap with a single loss term that free-runs the model during training exactly as evaluation will, decodes every rollout latent, and penalizes reconstruction error against ground truth. The term adds no parameters, costs training-time compute only, and reduces to the standard objective at weight zero, so every comparison in this paper is a one-flag A/B. On the chaotic Kuramoto-Sivashinsky equation, RDR raises valid prediction time (the time to first crossing of normalized error 0.5) from $3.87 \pm 0.23$ to $6.97 \pm 0.42$ time units at an identical 193,568 parameters, a $1.80\times$ improvement confirmed on seeds never used in selection and holding in 10 of 10 preregistered configurations at ratios of 1.71-2.50 $\times$ . The results come from a single system; a sweep in which the advantage grows with latent width is descriptive, and control experiments on two classic tasks are preliminary.

Figures & tables

Figure 1: Every preregistered configuration, canonical horizon. Each row is one configuration scored at three seeds; open circles are the posterior-only arm, filled circles the capacity-matched +RDR arm, with the per-row ratio at right. RDR wins 10 of 10, ratio 1.71–2.50

\times

. The fresh-seed confirmation row (R5) is the headline number; rows above it were selection rows.

Figure 2: Forecast error against ground truth on one held-out KS trajectory: two world models at an identical 193,568 parameters, free-running from the same initial condition. RDR holds lower error at every horizon (shaded); VPT is the first crossing of the 0.5 threshold,

t{=}3.9

for the posterior-only arm against

t{=}7.3

for RDR on this trajectory. Arm-level means over all 64 test trajectories at three fresh seeds are

3.87\pm 0.23

and

6.97\pm 0.42

tu.

model	params	VPT canonical (tu)	VPT long (tu)
latent + RDR, shared decoder (headline)	193,568	$6.97\pm 0.42$	$6.90\pm 0.44$
latent + RDR, split head (preregistered primary)	243,296	$6.47\pm 0.40$	$6.40\pm 0.40$
observation-space pushforward	230,464	$7.00\pm 0.26$	$6.90\pm 0.26$
latent, posterior-only	193,568	$3.87\pm 0.23$	$3.77\pm 0.23$

Table 1: Model classes at the winner’s training budget (320 epochs), canonical and long horizons, mean

\pm

sd over three seeds. The pushforward baseline sits between the two latent arms in parameter count. The headline arm matches it within 0.06 combined sd (root sum of squares of the two arms’ sds); the preregistered primary arm sits 1.11 combined sd below it.

Figure 3: The latent-size bracket at the matched recipe (decoder 512, 160 epochs,

\lambda{=}0.3

, final protocol, three paired seeds per rung). Left: VPT by arm; the posterior-only arm falls as the latent widens while the RDR arm rises, its top step within noise. Right: the per-rung paired ratio rises from 1.41

\times

to 2.50

\times

; the endpoints are separated by roughly 3.4 combined sd (sum of the endpoint sds), and the top step holds in 2 of 3 seeds. The bracket is descriptive: selection seeds, post hoc, no preregistered gate. Error bars are

\pm 1

sd over seeds.

Figure 4: Normalized error of the decoded free-running rollout against rollout depth, mean over the 64 held-out test trajectories, three fresh seeds per arm (seed spread shaded). The posterior-only arm crosses the 0.5 threshold at

3.87\,\mathrm{tu}

on average and the RDR arm at

6.97\,\mathrm{tu}

. Each seed’s crossing of its own curve is that seed’s reported VPT, so this figure and the headline number are one artifact. The persistence forecast is the dotted reference.

Figure 5: Distance between the free-running rollout latent and the posterior latent of the true state at the same time, against rollout depth: mean over the 64 held-out test trajectories, three fresh seeds per arm (seed spread shaded), normalized per model by its climatological latent scale so the two arms’ latent spaces are comparable. Both arms drift far off the posterior distribution and saturate near the climatological scale.

Figure 6: The reduced-data ladder on pendulum, mean return over 20 paired episodes, annotations giving the RDR-minus-posterior margin at each rung. Left: epochs held fixed, so reduced rungs receive proportionally fewer optimizer steps; RDR wins 20/20 paired episodes at every rung, holds near its full-data return down to the 25% rung, and the posterior arm degrades at every rung. Right: optimizer steps matched at every rung; the margins narrow to

+3

+13

-18

, the sign reverses at two cartpole rungs (6/20 paired wins at 25% and 50%), and at the shared scale the two arms coincide. The inset magnifies the marked sliver, across which the entire spread is about 30 return.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

config	posterior	+RDR shared	+RDR split, sharp (B)	+RDR split, ramp (C)
R1b: 320 epochs	$3.83\pm 0.21$	$7.13\pm 0.50$	6.53	$6.90\pm 0.50$
R5: winner, fresh seeds	$3.87\pm 0.23$	$6.97\pm 0.42$	6.13	$6.47\pm 0.40$
R1a: 240 epochs	3.37	$6.57\pm 0.71$	6.10	$6.37\pm 0.67$
R3b: $\lambda{=}0.6$	2.60	$5.80\pm 0.70$	5.10	$5.57\pm 0.76$
R3c: $\lambda{=}1.0$	2.60	$5.77\pm 0.15$	5.43	$6.07\pm 0.71$
R2b: latent 48	2.27	$5.67\pm 0.55$	5.47	$5.80\pm 1.35$

Appendix

Table 2: The preregistered sweep, all four arms, canonical horizon, VPT mean

\pm

sd (tu) over three seeds (sd suppressed where the archived summary reports arm means only). Long-horizon columns track canonical throughout (every arm within 0.15 tu of its canonical value) and are omitted for width; the winner’s long-horizon values appear in Section 4.2 . Rows sorted by shared-arm VPT. R5 is the winner configuration at fresh seeds 10–12; all other rows are selection rows at seeds 0–2.

config	seeds	posterior per-seed	+RDR shared per-seed
R1a: 240 epochs	0/1/2	3.4 / 2.9 / 3.8	6.7 / 7.2 / 5.8
R1b: 320 epochs	0/1/2	4.0 / 3.6 / 3.9	7.2 / 7.6 / 6.6
R2a: latent 24	0/1/2	2.7 / 2.8 / 2.8	4.9 / 4.7 / 4.6
R2b: latent 48	0/1/2	2.3 / 2.2 / 2.3	5.3 / 5.4 / 6.3
R3a: $\lambda{=}0.1$	0/1/2	2.7 / 2.3 / 2.8	5.9 / 5.4 / 4.6
R3b: $\lambda{=}0.6$	0/1/2	2.7 / 2.3 / 2.8	6.5 / 5.8 / 5.1

Appendix

Table 3: Per-seed VPT (tu), canonical horizon, for the two parameter-matched arms of every preregistered configuration. Every one of the 30 paired comparisons favors RDR. Per-seed values for the split-decoder arms and the long horizon are in the archived evaluation artifacts.

latent	posterior	+RDR shared	+RDR split, ramp (C)	ratio shared/posterior
16	$3.00\pm 0.36$	$4.20\pm 0.27$	$4.03\pm 0.21$	$1.41\pm 0.10$
24	$2.77\pm 0.06$	$4.73\pm 0.15$	$4.73\pm 0.40$	$1.71\pm 0.09$
32	$2.60\pm 0.27$	$5.90\pm 0.46$	$5.27\pm 0.40$	$2.30\pm 0.43$
48	$2.27\pm 0.06$	$5.67\pm 0.55$	$5.80\pm 1.35$	$2.50\pm 0.22$

Appendix

Table 4: The latent-size bracket (Figure 3 ), canonical horizon, matched recipe (decoder 512, 160 epochs,

\lambda{=}0.3

), three paired seeds per rung. Ratios are per-seed paired means. All four rungs are scored under the final protocol at the same recipe.

regime	data	pend. post / +RDR	$\Delta$ (wins)	cartp. post / +RDR	$\Delta$ (wins)
fixed epochs	10%	$-1400$ / $-1004$	$+396$ (20/20)	$-28$ / $+21$	$+48$ (20/20)
	25%	$-1263$ / $-191$	$+1072$ (20/20)	$-21$ / $+110$	$+132$ (20/20)
	50%	$-939$ / $-194$	$+745$ (20/20)	$+49$ / $+114$	$+64$ (20/20)
fixed steps	10%	$-177.0$ / $-194.9$	$-18.0$ (12/20)	90.3 / 95.6	$+5.4$ (17/20)
	25%	$-177.4$ / $-164.6$	$+12.8$ (12/20)	122.9 / 115.6	$-7.4$ (6/20)
	50%	$-179.6$ / $-176.2$	$+3.3$ (18/20)	129.2 / 124.8	$-4.3$ (6/20)

Appendix

Table 5: The reduced-data ladder, both regimes, both tasks: mean return of each arm with the RDR-minus-posterior paired margin and RDR’s paired win count over 20 episodes. Fixed-epoch rungs receive 480/240/60 optimizer steps as the data shrinks; step-matched rungs all receive the full-data 960 and were evaluated under the stateful planner, whose full-data rows appear in Table 6 . Rungs draw episode windows independently, so smaller rungs are not strict subsets of larger ones. Margins are paired, computed from unrounded per-episode returns, and can differ in the last digit from the difference of the rounded arm means.

row	planner	posterior	+RDR	$\Delta$ (wins)	post. gain	RDR gain
pendulum, seeds 0–19	per-step-reset	$-187.10$	$-179.27$	$+7.83$ (19/20)
	stateful	$-178.06$	$-177.45$	$+0.61$ (9/20)	$+9.04$	$+1.82$
pendulum, seeds 100–119	per-step-reset	$-263.02$	$-247.28$	$+15.74$ (18/20)
	stateful	$-244.90$	$-243.69$	$+1.22$ (13/20)	$+18.12$	$+3.59$
cartpole, seeds 0–19	per-step-reset	140.25	145.77	$+5.52$ (15/20)
	stateful	157.33	158.57	$+1.24$ (11/20)	$+17.08$	$+12.80$

Appendix

Table 6: The planner rollout comparison, all four measured rows: mean return per arm under the per-step-reset and stateful planner rollouts, with the paired margin and RDR’s win count over 20 episodes. The stateful rollout matches the one used in training. The last two columns give each arm’s gain from moving to the stateful planner: the posterior arm gains more in 4 of 4 rows. Margins are computed as in Table 5 .

quantity	posterior-only	+RDR (shared)
trainable parameters	193,568	193,568
batch size	64	64
optimizer steps (8 per epoch)	2,560	2,560
free-running rollout length $K$	128	128
decoder evaluations per optimizer step	20,416	28,608
decoder evaluations per training run	$52.3$ M	$73.1$ M

Appendix

Table 7: Compute accounting at the headline configuration, per training run on one NVIDIA H100. A decoder evaluation is one application of

D_{\psi}

to one latent vector. Both arms compute the

K{=}128

free-running rollout for

L_{\mathrm{R}}

; RDR additionally decodes each rollout latent, which is its entire marginal cost. Inference is identical for both arms.

tier	VPT (tu)	$\Lambda t$	status
persistence baseline	2.4	0.10
latent, posterior-only (winner config)	3.87	0.17	measured, fresh seeds
observation-space pushforward, small budget	5.77	0.25	the preregistered gate’s bar
latent + RDR, split head (primary arm)	6.47	0.28	fresh-seed confirmed
latent + RDR, shared decoder (headline)	6.97	0.30	fresh-seed confirmed
observation-space pushforward, matched budget	7.00	0.30	Section 4.3

Appendix

Table 8: Every arm on the absolute VPT scale for KS at

L{=}22

, with

\lambda_{\max}{=}0.043

[ 7 ] (one Lyapunov time = 23.26 tu). The published band and the reservoir row are different settings and carry the caveats below; neither is class-comparable to the latent arms.

Latent World Models

Latent world models learn compact predictive state representations to simulate how environments change and support planning, control, and reasoning.

See all 212 papers

Predictive but Not Plannable: RC-aux for Latent World Models
May 8, 2026Wenyuan Li, Guang Li, Keisuke Maeda +2Reachability Analysis World Model-Based Planning
Hokkaido University
Video World Models Need Predictable Latents
Oct 26, 2026Mattia Scardecchia, Sherry YangVideo World Models World Model Learning
New York University
Understanding the limits of latent world models
Oct 27, 2026Rodrigo Gonzalez Laiz, Steffen SchneiderWorld Models Latent World Models
Helmholtz Munich Munich Center for Machine Learning (MCML)

May 8, 2026J/K move · Enter open · S save

Predictive but Not Plannable: RC-aux for Latent World Models

Wenyuan Li, Guang Li, Keisuke Maeda, Takahiro Ogawa+1

Hokkaido University

Reachability Analysis World Model-Based Planning Model-Based Planning Latent World Models

Open paper AskPDF arXiv

A latent world model may achieve accurate short-horizon prediction while still inducing a latent space that is poorly aligned with planning. A key issue is spatiotemporal mismatch: these models are often trained with local predictive supervision, but deployed for long-horizon goal-directed search in latent spaces where Euclidean distance may not reflect what is reachable within a finite action budget. We present the Reachability-Correction auxiliary objective (RC-aux), a lightweight correction for this mismatch in reconstruction-free latent world models. RC-aux keeps the world-model backbone unchanged and adds planning-aligned supervision along two axes. Along the time axis, multi-horizon open-loop prediction trains the model beyond one-step consistency. Along the space axis, budget-conditioned reachability supervision, together with temporal hard negatives, encourages the latent space to distinguish states that are eventually reachable from those reachable within the current planning horizon. At test time, the learned reachability signal can also be used by a reachability-aware planner to favor trajectories that are both goal-directed and attainable under the available budget. We instantiate RC-aux on LeWorldModel and evaluate it under both continuation-training and matched-from-scratch settings. Across goal-conditioned pixel-control tasks and a LIBERO-Goal extension, RC-aux improves LeWM-style planning with modest additional cost. These results suggest that planning with latent world models depends not only on predictive accuracy, but also on whether the learned representation encodes the temporal and geometric structure required by downstream search. The code is available at https://github.com/Guang000/RC-aux.

Abstract

Latent World Models

Explore similar work

Predictive but Not Plannable: RC-aux for Latent World Models

Video World Models Need Predictable Latents

Understanding the limits of latent world models