Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning
This preprint introduces a training-time auxiliary loss for latent world models: a temporary linear head regresses the physical quantities used in a task's success criterion directly from latent states. On PushT and cube it raises planning success by 3.5 and 3.4 absolute points over the same baseline, with significant paired-test gains, while leaving test-time model and inputs unchanged. It matters because success criteria can tell world models exactly which physical quantities the latent state must retain for reliable planning.
Can supervising the physical quantities used in a task's success criterion—rather than arbitrary physical or proprioceptive quantities—make latent world-model planning more successful?
Latent world models rank action sequences by latent-state distance, but their training loss never requires the latent state to preserve the physical quantities used to judge success. In examined models, end-effector position readout error exceeds the success tolerance, so candidates that actually satisfy the goal cannot be reliably separated from failures.
Earlier latent world models either gave proprioceptive quantities as inputs, such as DINO-WM and V-JEPA 2-AC, which does not force later latent states to preserve them, or used designer-chosen physical quantities as auxiliary training targets, such as PhyLatent and PSG-JEPA, without tying the target set to the success criterion.
The method adds a shared linear head on both encoder output and predictor output that regresses the success-criterion quantities during training. Its squared error is added to the existing world-model loss, and the heads are removed afterward, so the planner, cost, and inputs stay unchanged at test time. The set of supervised quantities is fixed by the success criterion, not the designer, and the approach is evaluated on LeWM and a DINO-WM-type system with a frozen pretrained encoder.
Across four manipulation tasks with LeWM, supervision using all success-criterion quantities improved planning success: stack +5.05 points over a 30.1% baseline, cube +3.40 points over 68.5%, and PushT +3.50 points over 91.5%; on coffee, object-only supervision improved success by +3.90 points over 34.0%. Hand-position readout error fell from 3.2 cm in baseline LeWM to 0.14 cm with hand supervision, below the 2 cm success tolerance. In a frozen DINO-WM-type stack system, hand supervision improved success by +7.10 points with DINOv2-S and +3.60 points with DINOv2-L; supervising all criterion quantities improved success by +4.20 points with DINOv2-S. Evaluations used 2000 identical trial pairs for LeWM and 1000 for the frozen system, with paired two-sided McNemar tests. On stack, supervising only one quantity produced complementary gains and degradations, but supervising all criterion quantities never significantly degraded success in any reported task or trial stratum.
The authors note one training seed per condition, a weight sweep only for the hand on stack, and frozen-encoder replication limited to stack, with failure on frozen V-JEPA 2. The mechanism by which better readout translates to higher success is unresolved, and criteria defined by predicates such as contact or grasp are out of scope. The work is simulation-based across five tasks and one primary world model, with no real-robot validation; real-robot use would require reliable object-pose sensing. Code availability and compute cost are not reported in the paper.
This could improve learned world-model planners for vision-based robotic manipulation, especially in pick-and-place, assembly, or tending tasks where success is defined by end-effector and object pose tolerances. Robot OEMs and warehouse or industrial automation providers using learned planners could add this loss at training time at no inference cost. In the near term, it could be adopted in offline training pipelines for simulation-to-real manipulation models; production use would require reliable real-time estimation of success-criterion object poses and broader validation across tasks, seeds, and real robots.
The full text is not republished here because the paper's license does not allow it. Read the original on arXiv.