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

Jun 24, 2026cs.LG

Fast LeWorldModel

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.

Yuntian Gao, Xiangyu Xu

Xi’an Jiaotong University

Latent World Models World Models Joint-Embedding Predictive Architectures

Ask PDF

Sep 29, 2026cs.AI

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

Latent world models often struggle with long-horizon planning despite accurate short-term predictions. Recursive rollouts accumulate errors, while distance concentration in high-dimensional latent spaces can weaken goal discrimination. We introduce the Dual-Latent World Model (Dual-WM), which separates local execution and long-range planning through distinct state representations and dynamics models. The low-level model predicts action-conditioned transitions, while the high-level model uses learned macro-actions to plan over longer temporal spans. We also propose Long-Horizon Representation Learning with Weighted Rollout (LoRe), which supervises self-generated predictions at both levels. An analysis of recursive error propagation motivates exponential horizon weights with separate decay rates for the two temporal scales. During planning, the high-level model generates latent subgoals that the low-level model refines into actions for precise execution. We evaluate from-scratch Dual-WM on five goal-conditioned visual control tasks against the task-wise strongest baselines without actor-guided proposals. At goal offsets of 50 and 100 environment steps, mean success increases from 75.9% to 84.4% and from 61.4% to 69.5%, respectively. At offset 100, Dual-WM outperforms these baselines on all five tasks and improves mean success over LeWM by 30.8 percentage points. Ablations and supporting analyses provide evidence of more informative representations for goal evaluation and greater consistency under recursive prediction. These results highlight the value of separating temporal roles and training across multiple horizons for reliable latent planning. Our core implementation is available at https://github.com/DeLin1001/Dual-WM-Official.

Delin Zhao, Zhengrong Yue, Shaobin Zhuang +6

Nanjing University · Shenzhen University of Advanced Technology · Shanghai Jiao Tong University +1

Latent World Models Long-Horizon Planning Latent Space Planning

Ask PDF

Jun 30, 2026cs.AI

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

Learning visual world models for planning requires compact latent dynamics that remain sensitive to actions, yet reconstruction-free joint-embedding objectives can collapse to action-insensitive representations. We propose Delta-JEPA, an end-to-end reconstruction-free world model that augments latent forward prediction with a Latent Difference Action Decoder (LDAD). Unlike inverse decoders that infer actions from concatenated endpoint embeddings, LDAD reconstructs the executed action from the latent displacement between consecutive observations. This displacement-level supervision directly regularizes transition geometry: adjacent embeddings cannot collapse without losing action information, and different actions are encouraged to induce distinguishable latent changes for rollout-based planning. Delta-JEPA uses only latent prediction and action reconstruction, avoiding pixel reconstruction and distribution-matching regularizers. Across four visual continuous-control tasks, Delta-JEPA improves planning over JEPA-based and representation-learning world model baselines. Ablations show that displacement-based action decoding is consistently more effective than endpoint concatenation, and action-sensitivity analyses show clearer action-conditioned latent responses. These results indicate that supervising latent differences is a simple and effective mechanism for collapse-resistant and action-sensitive world model learning.

Zhenghao Zhang, Yuanxiang Wang, Zhenyu Guan +11

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

Latent World Models Video Joint Embedding Predictive Architecture S-Jepa

Ask PDF