Pub-AI: AI in Science

Latent Dynamics · 2026-09-30

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

Boyuan Zhang, Yingjun Du, Xiantong Zhen, Ling Shao

arXiv:2609.33595PDFCode

Auto-summarized: this summary was generated by a language model from the paper’s text and has not been reviewed by an editor. Check the paper before relying on it. Benchmark numbers appear only after a human has verified them.

TL;DR

The authors introduce SALT, an action-conditioned joint-embedding latent dynamics model with state-affine transitions and recursive rollout training. They argue one-step prediction error can fail to predict planning quality because planning composes transitions recursively, transforming introduced errors. Across four visual planning environments, SALT has higher one-step prediction error than LeWM but higher closed-loop success.

Why it matters

The paper isolates how prediction errors propagate during recursive latent planning. By connecting a structural condition (state-affine dynamics with state-independent Jacobians) to exact error propagation form, and by training with recursive rollout supervision, the authors show better planning can coincide with worse one-step prediction. This addresses a key mismatch between common one-step training objectives and how models are used in planning.

Method

  • Decompose multi-step rollout error into per-step errors and their propagation operators; show state-independent transition Jacobians correspond to affine-in-state differentiable dynamics.
  • Introduce SALT: action-conditioned state-affine latent transition where action modulates both the state transformation and an additive update; propagation depends only on the action sequence.
  • Train SALT with recursive multi-step rollout supervision by feeding each predicted latent state back into the transition (K=5) and applying SIGReg regularization (λ=0.09).

Limitation

The authors state their theoretical and empirical analyses address recursive error propagation, “one aspect of latent planning,” and “do not fully explain all factors determining planning performance.”

Abstract (from arXiv)

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%$.

Related papers