Beyond One-Step Accuracy: State-Affine Latent Transition for Reliable Visual Planning

Boyuan Zhang; Yingjun Du; Xiantong Zhen; Ling Shao

cs.ROSep 27, 2026

Beyond One-Step Accuracy: State-Affine Latent Transition for Reliable Visual Planning

Authors: Boyuan Zhang, Yingjun Du, Xiantong Zhen, Ling Shao

Organizations: UCAS-Terminus AI Lab, School of Engineering Science, University of Chinese Academy of Sciences · VIS Lab, University of Amsterdam · Central Research Institute, United Imaging Healthcare, Co., Ltd.

Abstract

Joint-embedding world models enable visual planning by learning action-conditioned dynamics in latent space. Yet they are commonly trained for one-step prediction on encoded states, while planning recursively applies the learned transition to its own predictions. One-step accuracy therefore does not capture how prediction errors propagate under recursive rollout. We decompose multi-step rollout error into the errors introduced at individual steps and their propagation through subsequent transitions. We show that state-affine dynamics are precisely the differentiable transitions with state-independent Jacobians, eliminating the nonlinear propagation residual and making the error propagation operators depend only on the action sequence. Guided by this result, we introduce SALT (State-Affine Latent Transition), an action-conditioned state-affine dynamics model in which the action modulates both the state transformation and the additive update. We train SALT through recursive multi-step rollout supervision, feeding each predicted latent state back into the transition so that training matches how the model is used during planning. Across four visual planning environments, SALT exhibits $1.48$ -- $2.19\times$ higher one-step prediction error than the matched LeWM baseline, yet improves closed-loop success in every environment by $10.0$ percentage points on average. On OGBench-Cube, the fraction of episodes that fail with a sharp rise in model-predicted cost after execution decreases from $23.3%$ to $2.0%$ .

Figures & tables

Panel of Figure 1 — Figure 1 : State dependence and current-state rollout reconstruction of the LeWM predictor. (a) Distribution of $\sigma_{\max}(\partial_{z}F)$ for the current-state Jacobian across 12 fixed action sequences while varying the current latent state. Points denote samples, the thick bar shows the interquartile range, and the dashed line marks $\sigma_{\max}=1$ . (b) Relative discrepancy between the rollout error and its current-state first-order reconstruction as the rollout horizon increases. The reconstruction uses only current-state Jacobians and omits propagation through earlier history inputs.

Method	Two-Room	Reacher	PushT	OGBench-Cube	Avg.
PLDM †	97	78	78	65	79.50
DINO-WM †	100	79	74	86	84.75
Sub-JEPA †	95.00	84.00	89.00	76.33	86.08
SD-JEPA †	90.00	88.00	97.30	72.00	86.83
SMWM †	99.00	66.00	83.00	84.00	83.00
Random	2.00 $\pm$ 2.00	12.00 $\pm$ 4.00	5.30 $\pm$ 2.30	42.00 $\pm$ 4.00	15.33

Table 1: Closed-loop planning success rates (%) across four environments. Controlled evaluations use

H=5

and report the mean

\pm

standard deviation over three evaluation seeds of 50 episodes each.

\dagger

quoted from the literature under different protocols. Bold marks the best mean among controlled rows.

			Training (it/s)
Predictor	Params	Forward (ms)	One-step	Rollout	Planning (s)
LeWM (ViT-S)	11.58M	2.36	6.5	2.9	35.2
SALT (Affine)	0.70M ( $\times$ 16.5)	0.29 ( $\times$ 8.1)	6.0	2.4	11.4 ( $\times$ 3.1)

Table 2: Predictor complexity and computational efficiency. Parameter counts and single-step forward times cover the predictor together with its projection head. Planning time is the CEM wall-clock time averaged over the four environments and

H\in\{5,10,15,20\}

. Additional timing details are provided in Appendix B.6 .

Panel of Figure 2 — Figure 2 : From one-step accuracy to recursive error propagation. (a) Relative one-step prediction error and closed-loop success across four environments. (b) Recursive rollout error on PushT. (c) Spectral norm of the measured propagation operator on PushT; the band spans the minimum and maximum over nine measurements.

Objective	LeWM	SALT
One-step	82.7	84.4
Rollout	77.8	92.7
$\Delta$	-4.9	+8.3

Table 3: Transition structure and training objective. Mean closed-loop success rate (%) over Two-Room, Reacher, PushT, and OGBench-Cube, with equal weight assigned to each environment. Training windows, target frames, and training budgets are matched across objectives.

Figure 3 : Failure anatomy on OGBench-Cube. Among failed episodes, mismatch failures satisfy $r=c/\hat{c}>5$ , while unflagged failures satisfy $r\leq 5$ . Here, $\hat{c}$ and $c$ are model-predicted terminal costs before execution and during subsequent replanning, respectively. The examples are reproduced LeWM failures, and all images are real observations. Bars report counts over $150$ paired episodes per model.

Figure 4 : Long-horizon planning performance. Success rate on the four environments at $H\in\{10,15,20\}$ . Arrows mark the performance difference between SALT and LeWM at $H=10$ and $H=20$ .

Method	Success rate
Random	32.7 $\pm$ 4.2
LeWM	36.0 $\pm$ 5.3
SALT	54.0 $\pm$ 3.5

Table 4: Success rate (%) on Cube-Double.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Component	Symbol	Shape	Parameters	Initialization
Generator of $U$	$W_{u}$	$192\times 192$	$36{,}864$	$0$
Generator of $V$	$W_{v}$	$192\times 192$	$36{,}864$	$0$
Singular-value parameters	$s$	$192$	$192$	$3.0$
Gate	$W_{g}$	$16\times 192$	$3{,}072$	$\mathcal{U}(-1/\sqrt{192},\,1/\sqrt{192})$
Modulation modes	$N_{1},\dots,N_{16}$	$16\times 192\times 192$	$589{,}824$	$0$
Action injection	$B$	$192\times 192$	$36{,}864$	$0$

Appendix

Table 5 : Parameters of the SALT transition model (

d=192

R=16

). Counts are stored parameters. Only the skew-symmetric parts of

W_{u}

and

W_{v}

affect the transition, corresponding to

18{,}336

independent entries for each matrix.

Figure 5 : Decoded open-loop rollouts across four environments. The first three columns ( $t=0,5,10$ ) are ground-truth context frames from the common evaluation interface. Only the final frame initializes SALT , which predicts seven subsequent latent states using dataset actions. At frameskip $5$ , these predictions cover $35$ additional environment steps ( $t=15,\ldots,45$ ). Images are produced by a post-hoc decoder trained only on ground-truth latents and excluded from world-model training and planning.

Figure 6 : Rollout error and propagation across four environments. Top: relative rollout error as a function of rollout step $k$ . Bottom: the spectral norm of the measured propagation operator $\Phi_{0\to k}$ . Results use test seeds $1$ , $2$ , and $3$ . Blue circles denote reproduced LeWM; orange squares denote SALT . Each model step corresponds to five environment steps.

Model	Mismatch	Unflagged
LeWM	$35/150$ ( $23.3\%$ )	$17/150$ ( $11.3\%$ )
SALT	$3/150$ ( $2.0\%$ )	$6/150$ ( $4.0\%$ )

Appendix

Table 6 : Episode-level failure rates by type on OGBench-Cube. Rates are computed over

150

paired episodes per model using the fixed threshold

r=5

. LeWM denotes the reproduced baseline.

Environment	Horizon	LeWM	SALT	Gap (pp)
Two-Room	$10$	$63$	$74$	$\mathbf{+11}$
	$15$	$54$	$65$	$\mathbf{+11}$
	$20$	$47$	$58$	$\mathbf{+11}$
Reacher	$10$	$26$	$30$	$\mathbf{+4}$
	$15$	$20$	$26$	$\mathbf{+6}$
	$20$	$18$	$26$	$\mathbf{+8}$

Appendix

Table 7 : Success rate (%) under the long-horizon protocol. Results correspond to the long-horizon experiment reported in the main text.

Configuration	Planning success
Action-independent transition matrix	$11.3\pm 3.1$
Full structure	$90.7\pm 3.1$

Appendix

Table 8 : Effect of removing the action-dependent modulation of the transition matrix. Success rate (%) on Reacher over three evaluation seeds.

Figure 7 : Success rate against the SIGReg weight on PushT. Shading shows standard deviation across three evaluation seeds; the dashed line marks the default $\lambda=0.09$ .

Figure 8 : Success rate of SALT against latent dimension on Two-Room and OGBench-Cube. The horizontal axis uses a base-two log scale; error bars show standard deviation across three evaluation seeds.

Explore similar work

Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent World Models World Models
Xi’an Jiaotong University
Beyond a single latent space: a dual-latent world model for long-horizon planning
Sep 29, 2026Delin Zhao, Zhengrong Yue, Shaobin Zhuang +6Latent World Models Long-Horizon Planning
Nanjing University · Shenzhen University of Advanced Technology · Shanghai Jiao Tong University +1
Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding
Jun 30, 2026Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan +11Latent World Models Video Joint Embedding Predictive Architecture
School of Computer Science and Technology, University of Chinese Academy of Sciences, Beijing · Institute of Information Engineering, Chinese Academy of Sciences, Beijing · School of Computer Science and Technology, Harbin Institute of Technology, Weihai +2

Jun 24, 2026, cs.LGJ/K move · Enter open · S save

Fast LeWorldModel

Yuntian Gao, Xiangyu Xu

Xi’an Jiaotong University

Latent World Models World Models Joint-Embedding Predictive Architectures

Open paper AskPDF arXiv

Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.

Abstract

Explore similar work

Fast LeWorldModel

Beyond a single latent space: a dual-latent world model for long-horizon planning

Delta-JEPA: Learning Action-Sensitive World Models via Latent Difference Decoding