Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning
Latent world models plan by scoring action sequences with latent-space distances, yet physical success criteria demand precision finer than these models typically encode. This work proposes a training-only auxiliary loss that supervises encoder and predictor outputs against the exact success-criterion quantities via a shared linear head, then discards the head at test time so planning remains unchanged. On manipulation tasks the approach yields statistically significant absolute gains of 3-5%, with the key insight that a task's success criterion itself should dictate what the latent state must retain.
The paper presents a simple, principled intervention: add a linear regression head on $\bm{z}_t$ and $\hat{\bm{z}}_t$ to regress the success-criterion quantities $\bm{p}_t^{(k)}$ during training, weighted by $\lambda_k$ and discarded at test time. The effect size is modest but statistically significant, and the stratum-level analysis reveals that partial supervision can improve one trial class while degrading another. However, the experimental footprint is thin—only one training seed per main condition, frozen-encoder replication limited to a single task, and the auxiliary-loss weight was swept only for the hand in stack. The authors are commendably transparent about these limits, noting that the mechanism linking improved readout to higher success rate remains unidentified and that the frozen V-JEPA 2 system shows improved readout yet lower success.
The central empirical finding is robust within the tested scope: in all four manipulation tasks, supervising the full set of success-criterion quantities raises the success rate without ever causing significant degradation overall or per stratum. The ablation over trial strata is a genuine strength, showing that on stack, hand-only supervision improves reach-only trials (+5.92%) while degrading manipulation trials (-8.24%), an effect canceled out at the aggregate level. The authors also correctly demonstrate that feeding a quantity as an input does not guarantee the latent state retains it at criterion precision, whereas using it as a training target does.
The main limitation is experimental breadth. Each reported comparison rests on a single training seed, which the authors acknowledge under Limitations. The replication on frozen-encoder systems is restricted to stack, and explicitly fails for frozen V-JEPA 2, undercutting the universality of the proposed rule. Furthermore, the mechanism of complementary degradation when supervising only a subset of quantities is not explained by readout accuracy of the unsupervised quantities, and the authors admit that the route by which the auxiliary loss raises the success rate is not identified.
A second concern is that the loss weight was only swept for the hand in stack; other tasks used default or stack-derived values without task-specific tuning, leaving open whether the gains are sensitive to hyperparameters.
The evidence supports the core claim that criterion-aligned supervision improves planning success rates on LeWM, with gains replicated on a DINO-WM–type frozen encoder for stack. The comparison with prior work is precise: the authors distinguish their target-based supervision from DINO-WM and V-JEPA 2–AC, which concatenate proprioceptive inputs but do not supervise the latent state, and from PhyLatent and PSG-JEPA, which supervise physical quantities chosen by the designer rather than aligned with the task-specific success criterion. They also appropriately caution against conflating proxy metrics with downstream success, supported by the frozen V-JEPA 2 counterexample.
The paper provides detailed hyperparameters, model architectures, and dataset sources in Appendix A, including the auxiliary-loss weights listed in Table 4 and the exact CEM settings. The evaluation protocol—paired McNemar tests on identical trial pairs drawn from demonstrations—is clearly described. However, no code or data release is mentioned, and the reliance on a single training seed per main condition means that variance across runs is uncharacterized. The five additional stack models in Appendix D suggest the aggregate result is directionally consistent, but the range of overall improvements (+1.35% to +11.35%) indicates substantial seed-dependent variability that a single run cannot capture.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.