Latent Dynamics · 2026-10-02
Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning
Takumi Hara, Kanata Suzuki
AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.
TL;DR
The authors show that latent world-model planning can fail because the latent distance used for ranking does not reliably preserve the physical quantities that define success. They propose an auxiliary training loss that regresses the success-criterion quantities from both encoder and predictor latent outputs using a linear head (discarded at test time). This improves success rate on PushT and cube by +3.5% and +3.4% (absolute) and is statistically significant.
Why it matters
It directly aligns what latent state must retain with what the task judges as success, without changing test-time planning or inputs. The method specifies supervision from the benchmark’s success criterion and adds a training-time regression head whose only purpose is to make success-criterion quantities linearly decodable with sufficient precision, improving downstream success rates.
Method
- Train a latent world model with an auxiliary loss: a linear head regresses success-criterion quantities from both encoder output z_t and predictor output ẑ_{t+1}, adding regression error to the training loss.
- At test time, discard the linear heads so the model, computation, and latent-distance cost for planning remain unchanged.
- Evaluate via paired two-sided McNemar tests on identical trial pairs and study which supervised quantities affect success by task and trial strata.
Limitation
Reachability issues are out of scope: the authors note they do not cover tasks where reachability decides success, such as mazes or tasks with walls.
Abstract (from arXiv)
Latent world models plan by scoring candidate action sequences with distances in latent space. However, task success is judged by physical quantities, which we call the success-criterion quantities. In all four latent world models we examine, the end-effector position is encoded in the latent state with an error larger than the success criterion allows. Such a latent state cannot separate successful candidates from failing ones. We propose an auxiliary loss that uses success-criterion quantities as training targets, whereas existing latent world models take them only as inputs. During training, a linear head on the encoder and predictor outputs regresses the success-criterion quantities, and the regression error is added to the training loss. The head is discarded after training, so the model, its cost, and its inputs at test time are unchanged. This loss alone improves the success rate on PushT and cube by 3.5% and 3.4% (absolute), respectively, and both improvements are statistically significant. A success criterion thus specifies what a world model must retain in its latent state, and we show that it can serve directly as a training target.
Related papers
- When Low Prediction Error Misleads Planning: Diagnosing Representation, Dynamics, and Decision Failures in Latent World Models
- Action-Consequence Alignment for Reliable Planning and Self-Improving in Latent World Models
- One-Step Prediction Fidelity Is a Poor Proxy for Planning Performance in JEPA World Models
- Knowing When to Stop Thinking: A Dissection of Expected Free Energy as a Test-Time Compute Controller for Latent World-Model Planning
- CausalDreamer: Learning Predictive World Models with Latent Disentanglement
Publication record
Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.
First public release: 2026-10-01 · Indexed by Pub-AI: 2026-10-02 UTC
Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning
Preprint posted: 2026-10-01 · Original record ↗
- arxiv
- 2610.01224
Explicit source versions (1)
- v1 · 2026-10-01
- Preprint available
Source date: 2026-10-01 · First detected: 2026-10-25 UTC · Historical baseline
Source evidence (1)
- Official evidence ↗ · atom:published (official metadata)
Abstract · arxiv
Latent world models plan by scoring candidate action sequences with distances in latent space. However, task success is judged by physical quantities, which we call the success-criterion quantities. In all four latent world models we examine, the end-effector position is encoded in the latent state with an error larger than the success criterion allows. Such a latent state cannot separate successful candidates from failing ones. We propose an auxiliary loss that uses success-criterion quantities as training targets, whereas existing latent world models take them only as inputs. During training, a linear head on the encoder and predictor outputs regresses the success-criterion quantities, and the regression error is added to the training loss. The head is discarded after training, so the model, its cost, and its inputs at test time are unchanged. This loss alone improves the success rate on PushT and cube by 3.5% and 3.4% (absolute), respectively, and both improvements are statistically significant. A success criterion thus specifies what a world model must retain in its latent state, and we show that it can serve directly as a training target.