PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Boyuan Zhang, Yingjun Du, Xiantong Zhen, Ling Shao

arXiv:2609.33595v1cs.ROcs.LG

TL;DR

One-step prediction accuracy does not capture how errors propagate when latent world models are recursively applied during planning. The paper introduces SALT, an action-conditioned state-affine transition trained with recursive rollout supervision. Across four visual planning environments, SALT achieves higher closed-loop success despite larger one-step error than matched LeWM.

  • Problem

    One-step training on encoded states does not constrain how prediction errors propagate when a latent transition is recursively applied during planning.

    Source: §1 · p. 1

  • Method

    SALT uses action-conditioned state-affine dynamics and recursive multi-step rollout supervision that feeds predicted latent states back into the transition.

    Source: §1 · p. 2

  • Results

    Across four environments, SALT has 1.48–2.19× the one-step prediction error of matched LeWM yet achieves higher mean closed-loop success in every environment.

    Source: §1 · p. 2

  • Takeaways & Limitations

    In the evaluated setting, one-step prediction error alone is not a reliable proxy for closed-loop planning quality.

    Source: §4.3 · p. 8

    The shared base transition matrix is constrained to have σmax(A0) < 1.

    Source: §3.4 · p. 6

Abstract

from arXiv · show

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

1. Introduction

One-step training of latent world models does not constrain how prediction errors propagate when transitions are recursively applied during planning. The paper characterizes state-affine dynamics and introduces SALT, whose larger one-step error accompanies higher closed-loop success across four environments.

Introduction

  • One-step prediction objectives supervise transitions from encoded states, while planning recursively applies the model to its own predictions.Rollout error reflects both errors introduced at individual steps and their propagation through later transitions.Source: §1 · pp. 1–2
  • For differentiable transitions, state-affine dynamics are precisely those with state-independent Jacobians, eliminating the nonlinear propagation residual.This makes propagation operators depend on the action sequence, separating individual prediction accuracy from error transformation during rollout.Source: §1 · p. 2
  • SALT makes the transition affine in latent state, lets actions modulate both the state transformation and additive update, and trains with recursive multi-step supervision.Predicted latent states are fed back into the transition during training to match recursive use in planning.Source: §1 · p. 2
  • Across four environments, SALT has 1.48–2.19× the one-step prediction error of matched LeWM yet achieves higher mean closed-loop success in every environment.Its average success is 92.7%, 10.0 percentage points above the baseline; on OGBench-Cube, mismatch failures decrease from 23.3% to 2.0%.Source: §1 · p. 2

2. Related Work

Prior work develops latent world models, joint-embedding representations, multi-step learning, and structured dynamics for prediction and planning. This paper focuses on the transition structure that makes error propagation state-independent.

Latent World Models for Visual Planning

Stable and Structured Joint-Embedding World Models

  • Joint-embedding work has improved representation learning, while prior work also studies autoregressive error growth; these approaches do not enforce a state-independent transition Jacobian (Assran et al., 2023; Balestriero and LeCun, 2025; Maes et al., 2026).This paper focuses on the structural condition yielding an exact error decomposition for action-conditioned state-affine transitions.Source: §2 · p. 3

Multi-Step Prediction and Structured Latent Dynamics

3. Method

The method frames visual planning as recursive action-conditioned prediction in latent space, where rollout errors combine newly introduced errors with their propagation. It characterizes state-affine transitions, parameterizes action-conditioned dynamics, and trains the model with recursive rollout supervision.

3.1. Problem Setup

  • Planning predicts latent states by recursively applying the transition over candidate action sequences, whereas one-step training supervises transitions from encoded states.A rollout's error depends on both per-step prediction errors and how later transitions propagate them.Source: §3.1 · pp. 3–4

3.2. Recursive Error Propagation

  • The rollout decomposition separates errors introduced by individual transitions from their propagation through later transitions, where each later transition’s state Jacobian transforms earlier errors and residuals capture terms beyond first-order linearization.Source: §3.2 · p. 4

3.3. State-Affine Dynamics

  • A differentiable transition has a state-independent Jacobian if and only if it is affine in the latent state; then the nonlinear residual vanishes and propagation matrices follow from the action conditions.For a fixed action sequence, those matrices can be computed without generating the latent rollout.Source: §3.3 · pp. 5–6

3.4. Action-Conditioned State-Affine Transition

  • SALT parameterizes an action-conditioned affine transition so actions can change the state transformation while preserving affine dependence on latent state.The base transition matrix is constrained to have σmax(A0) < 1.Source: §3.4 · p. 6

3.5. Recursive Rollout Training

  • Recursive rollout training feeds predicted latents back into the transition and supervises each predicted state, aligning training with recursive planning use.The method uses K = 5 predictions from a six-frame training window and combines rollout supervision with SIGReg.Source: §3.5 · p. 6

4. Experiments

Across four visual planning environments, SALT improves closed-loop success over reproduced LeWM despite higher one-step prediction error. The experiments examine planning efficiency, recursive error behavior, failure types, longer horizons, and a two-object task.

4.1. Experimental Setup

  • Evaluations use four visual planning environments, a matched encoder and planner for SALT and reproduced LeWM, and three seeds of 50 episodes per environment.Planning uses CEM with horizon and replanning interval both set to five; goals are observations 25 environment steps after the initial state in the same offline trajectory.Source: §4.1 · p. 7

4.2. Closed-Loop Planning Performance

  • SALT reaches 92.7% average success, 10.0 percentage points above reproduced LeWM, and has higher mean success in all four environments.On OGBench-Cube, success rises from 67.3% to 94.0%, a 26.7-point gain.Source: §4.2 · p. 7, Table 1 · p. 7
  • SALT uses 0.70 M parameters versus LeWM’s 11.58 M, with 8.1× faster single-step prediction and 3.1× faster CEM planning.Timing comparisons cover the predictor and projection head where present; CEM time is averaged across four environments and horizons 5, 10, 15, and 20.Source: §4.2 · pp. 7–8, Table 2 · p. 8

4.3. Why One-Step Accuracy Fails to Predict Planning Quality

  • Across four environments, SALT’s one-step prediction error is 1.48×–2.19× LeWM’s, yet its closed-loop success is higher in every case.On PushT, rollout-error rankings reverse over time, while measured propagation-operator norms grow by 4.28× for SALT and 8.74× for LeWM over twenty steps.Source: §4.3 · p. 8, Figure 2 · p. 9
  • With matched training budgets, rollout training raises SALT’s mean success from 84.4% to 92.7% but lowers LeWM’s from 82.7% to 77.8%.The full configuration achieves the highest mean, indicating that the benefit of recursive rollout supervision depends on transition structure.Source: §4.3 · pp. 8–9, Table 3 · p. 9

4.4. Planning Reliability: Failure Anatomy on OGBench-Cube

  • On OGBench-Cube, mismatch failures fall from 23.3% of paired episodes for LeWM to 2.0% for SALT.Unflagged failures fall from 11.3% to 4.0%; mismatch failures are defined by a greater-than-five ratio of post-execution to pre-execution model-predicted terminal costs.Source: §4.4 · p. 9, Figure 3 · p. 10

4.5. Long-Horizon Stress Test

  • SALT achieves higher success in all twelve environment–horizon combinations at horizons 10, 15, and 20.At horizon 20, its advantage is 11, 8, 4, and 17 percentage points on Two-Room, Reacher, PushT, and OGBench-Cube, respectively.Source: §4.5 · p. 10, Figure 4 · p. 10

4.6. Planning with Multiple Objects

  • On Cube-Double, SALT reaches 54.0% success, compared with 36.0% for LeWM and 32.7% for random.The task requires two cubes to reach their target poses simultaneously.Source: §4.6 · p. 10, Table 4 · p. 11

5. Conclusion

The conclusion links state-affine dynamics and recursive rollout training to improved closed-loop planning despite larger one-step errors. Across four environments, SALT also shows milder measured propagation and fewer OGBench-Cube failures marked by sharply higher post-execution predicted costs.

5. Conclusion

  • Across four environments, SALT achieves higher mean closed-loop success than reproduced LeWM despite larger one-step prediction error.The conclusion also reports milder measured propagation-operator amplification over longer rollouts.Source: §5 · p. 11
  • State-affine dynamics have state-independent Jacobians, and SALT combines this structure with action conditioning and recursive rollout training.The analysis separates per-transition errors from their subsequent propagation.Source: §5 · p. 11
  • On OGBench-Cube, SALT reduces failures characterized by sharply higher model-predicted costs after execution and re-observation.Source: §5 · p. 11

A.1. Proof of the Multi-Step Error Decomposition

The proof establishes a multi-step error expansion into introduced errors and their subsequent propagation, then proves that state-independent Jacobians characterize state-affine transitions.

A.1. Proof of the Multi-Step Error Decomposition

  • Equation 13 expresses the rollout error as a sum of per-step errors and residuals, each transformed by the subsequent Jacobian product.The proof verifies the expansion by induction, with the one-step case satisfying r1 = 0 and δ1 = ε1.Source: §A.1 · p. 14
  • The residual bound assumes the state derivative is L-Lipschitz, with the same constant for every action condition.The proof applies the fundamental theorem of calculus along the segment from zt−1 to zt−1 + δt−1.Source: §A.1 · p. 14

A.2. Proof of the State-Affine Characterization

  • If the state Jacobian is constant for a fixed action condition, integrating along the segment between states yields F(z, c) = A(c)z + B(c).Conversely, differentiating this affine form gives a state-independent Jacobian.Source: §A.2 · pp. 14–15

B.1. Environments and Datasets

The appendix specifies SALT’s datasets, model and recursive training, evaluation protocol, and diagnostics. It also describes how the rollout and propagation measurements are constructed and qualified.

B.1. Environments and Datasets

  • Experiments follow LeWM benchmark settings across Two-Room, Reacher, PushT, and OGBench-Cube, with dataset windows split into 90% training and 10% validation.Evaluation initial states and goals are sampled from the offline dataset.Source: §B.1 · p. 15

B.2. Training Configuration

  • The visual encoder is a scratch-trained ViT-Tiny, and the training configuration uses 10 epochs, batch size 128, and latent dimension 192.The listed optimizer settings include AdamW with learning rate 5 × 10−5 and weight decay 10−3.Source: §B.2 · p. 15

B.3. Transition Model

  • SALT acts directly in regularized latent space, with action-conditioned state transformation and additive update; its base matrix has singular values below one for finite parameters.The parameter table reports d = 192 and R = 16; the singular-value constraint applies to the base matrix, not the full action-conditioned matrix.Source: §B.3 · pp. 15–16, Table 5 · p. 16

B.4. Training Procedure

  • Recursive training initializes from the first latent and supervises all five predicted steps equally, while SIGReg uses all six encoded frames and target latents are not detached.The sampled window contains six observations at five-environment-step intervals and spans 25 environment steps.Source: §B.4 · p. 16, §B.7 · p. 17

B.5. Planner and Evaluation Protocol; B.7. Baseline Provenance

  • Evaluation uses CEM with a five-step planning horizon and replanning interval; matched models share encoders, action encoders, latent dimensions, optimizer settings, and planners.LeWM retains its Transformer predictor and one-step objective, while SALT uses the rollout configuration; evaluation retains LeWM’s three-frame history interface.Source: §B.5 · p. 16, §B.7 · p. 16

C.1. Measurement Protocol; C.3. Propagation-Operator Growth

  • For LeWM propagation measurements, Jacobians differentiate only the most recent latent input while earlier history frames are fixed, so they are not the full augmented-state Jacobian.The PushT diagnostic aggregates nine window–batch combinations and reports geometric-mean operator norms with minimum-to-maximum bands.Source: §C.1 · p. 18, §C.3 · p. 18

C.2. Rollout Reconstruction; C.3. Propagation-Operator Growth; C.4. Control for the Spectral Constraint

  • SALT’s rollout-reconstruction discrepancy stays at 10−6 across horizons 1–20, while reproduced LeWM reaches 49.4% at H = 20; SALT’s propagation growth is also lower in the reported comparison.The measured twenty-step growth changes from 4.28× to 4.42× when SALT’s base-matrix constraint is removed, versus 8.74× for reproduced LeWM; these measurements do not establish global contraction.Source: §C.2 · p. 18, §C.3 · p. 18, §C.4 · p. 19

C.5. Decoded Open-Loop Rollouts

  • Decoded open-loop examples retain recognizable agent, arm, or object configurations, but the post-hoc decoder is used only for visualization, not training, planning, or performance measurement.Figure 5 uses dataset actions and shows predictions spanning 35 additional environment steps after the final context frame.Source: §C.5 · p. 19, Figure 5 · p. 20

C.6. Rollout Error and Propagation Across Four Environments

Across four environments, SALT has larger initial prediction error but generally lower measured propagation norms at intermediate and later rollout steps. The passages distinguish these operator measurements from realized rollout error.

C.6. Rollout Error and Propagation Across Four Environments

  • Across four environments, SALT has larger initial prediction error but generally lower measured propagation norms at intermediate and later steps.Figure 6 reports test seeds 1, 2, and 3, with each model step corresponding to five environment steps.Source: §C.6 · p. 21, Figure 6 · p. 21
  • Smaller propagation norms do not imply smaller rollout errors at every horizon; on OGBench-Cube, final-step errors are not lower for SALT.On Reacher, SALT’s norm decreases after an initial rise; on PushT, it continues to grow more slowly than LeWM’s.Source: §C.6 · p. 21

D.1. Paired Evaluation and Failure Criterion

The paired evaluation fixes the failure threshold across models and seeds, then partitions OGBench-Cube failures by post-execution cost changes. The figure examples show two LeWM failures from seed 1.

D.1. Paired Evaluation and Failure Criterion

  • r = c/ĉ classifies failed episodes as mismatch failures when r > 5 and unflagged failures when r ≤ 5.The threshold of 5 was chosen within a gap in LeWM's seed-1 ratio distribution and held fixed for both models and remaining seeds.Source: §D.1 · pp. 21–22

D.2. Failure Composition

  • 94.0% for SALT versus 65.3% for LeWM, with 43 episodes won only by SALT and none won only by LeWM.Both models succeed on 98 paired episodes, and neither succeeds on 9.Source: §D.2 · p. 22, Table 6 · p. 22

D.3. Figure Details

  • Both Figure 3 examples are LeWM failures from seed 1, with ratios near the seed-1 median for their respective failure types.The start and goal images are dataset observations, and the post-execution image is recorded after the first planned sequence; all are real observations.Source: §D.3 · p. 22

E.1. Evaluation Protocol

The long-horizon evaluation varies planning horizon and replanning interval across four environments. SALT exceeds LeWM in all twelve environment–horizon combinations.

E.1. Evaluation Protocol

  • The stress test uses a 100-step target offset, a 300-step interaction budget, and horizons and replanning intervals H ∈ {10, 15, 20}.With frameskip 5, full-sequence execution intervals are 50, 75, and 100 environment steps; other settings follow Section B.5, and LeWM is the reproduced baseline.Source: §E.1 · p. 22

E.2. Results Across Four Environments

  • SALT exceeds LeWM in all twelve environment–horizon combinations, with success-rate gaps of 4–25 percentage points.Table 7 reports the success rates underlying Figure 4.Source: §E.2 · p. 22, Table 7 · p. 23

F. Ablations and Robustness

The ablations test action-dependent transition modulation, SIGReg weight, and latent dimensionality. Comparisons are valid within each evaluation round, not across ablation groups.

F. Ablations and Robustness

  • Absolute results across different ablation groups should not be interpreted as directly comparable.Each ablation compares variants within the same evaluation round.Source: F. Ablations and Robustness · p. 23

F.1. Action-Conditioned Transition

  • Removing action-dependent modulation of the transition matrix substantially reduces Reacher success; the additive action term alone does not match the full model in this setting.The ablation retains action inputs through the additive term but replaces A(c) with shared A0.Source: §F.1 · p. 23, Table 8 · p. 23

F.2.1. Regularizer Weight

  • PushT performance remains stable near the default SIGReg weight λ = 0.09 and decreases toward the tested extremes.Figure 7 plots success rate against SIGReg weight, with standard-deviation shading across three evaluation seeds.Source: §F.2.1 · p. 24, Figure 7 · p. 24

F.2.2. Latent Dimensionality

  • Across the tested fourfold latent-dimension range, success varies by at most approximately three percentage points in each environment.The visual backbone remains fixed while projection-head output, action-condition dimension, and transition parameters are adjusted to match d.Source: §F.2.2 · p. 24, Figure 8 · p. 24

G. Limitations and Scope

The paper limits its scope to JEPA-style world models using recursive latent prediction and addresses recursive error propagation as one factor in planning performance.

G. Limitations and Scope

  • The analysis studies JEPA-style world models that plan through recursive latent prediction, not architectures with different prediction and control interfaces such as vision-language-action models.Source: G. Limitations and Scope · p. 24
  • The theoretical and empirical analyses address recursive error propagation as one aspect of latent planning and do not fully explain all factors determining planning performance.Source: G. Limitations and Scope · p. 24

Related papers

Loading 2609.33595v1…