PaperlayerView exact revision on arXiv

Source-linked AI summary

Supervise What Decides Success: Criterion-Aligned Auxiliary Losses for Latent World-Model Planning

Takumi Hara, Kanata Suzuki

arXiv:2610.01224v1cs.LGcs.RO

TL;DR

Latent world models rank action sequences by latent distance, while task success is judged by physical criterion quantities that training does not require the latent state to retain. The paper adds auxiliary regression supervision on those quantities during training, then discards the head. Across four manipulation tasks, this supervision significantly improves success over baseline, with replication using a frozen pretrained encoder.

  • Problem

    Latent planning ranks candidates by latent distance, but training does not supervise success-criterion quantities, and hand-position readout error exceeds the criterion tolerance in all four examined models.

    Source: §1 · p. 1

  • Method

    A training-only auxiliary loss regresses success-criterion quantities from encoder and predictor outputs, with the head discarded so test-time planning remains unchanged.

    Source: §1 · p. 2

  • Results

    Across four manipulation tasks, criterion-quantity supervision significantly improves success over baseline, with replication on a frozen pretrained encoder.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The success criterion specifies what a world model must retain in its latent state and can serve directly as a training target.

    Source: §1 · p. 2

    Each condition has a single training seed, so comparisons rely on paired tests over identical trial pairs.

    Source: §5 · p. 9

Abstract

from arXiv · show

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.

1 INTRODUCTION

Latent planners rank actions by latent distance, yet success is judged by physical quantities that training does not explicitly require the latent state to retain. The paper adds criterion-targeted auxiliary supervision and reports higher success across four manipulation tasks.

1 INTRODUCTION

  • Across all four examined latent world models, hand-position readout error exceeds criterion tolerance and does not vanish with representation size or pretraining.Because latent distance alone ranks candidates, planning requires criterion-relevant physical quantities to be retained with sufficient precision.Source: §1 · p. 1
  • Prior approaches either feed physical quantities as inputs without requiring their retention after prediction, or use designer-selected targets not tied to the success criterion.The cited examples are DINO-WM and V-JEPA 2-AC for inputs, and PhyLatent and PSG-JEPA for targets.Source: §1 · pp. 1–2
  • A training-only auxiliary loss regresses criterion quantities from encoder and predictor outputs; its head is discarded, leaving test-time inputs, model, and planning cost unchanged.The success criterion, rather than designer choice, fixes which quantities are supervised.Source: §1 · p. 2
  • Lower readout error alone need not improve planning success, so the paper evaluates the auxiliary loss by success rate and ablates which quantities are supervised.The frozen V-JEPA 2 experiment combines corrected readout with a lower success rate.Source: §1 · p. 2
  • On four manipulation tasks, supervising success-criterion quantities significantly improves success rate over baseline, with replication on a frozen pretrained encoder.Supervising all quantities in the criterion never significantly degrades success on any task, for scratch-trained and frozen-encoder models.Source: §1 · p. 2

2 METHOD

Planning ranks candidate action sequences by predicted-goal latent distance alone, so the latent representation must retain criterion quantities at the precision needed near the goal. The method adds training-only regression heads on encoder and predictor states without changing inference.

2.1 PLANNING IN LATENT SPACE

  • For each candidate, the planner predicts a terminal latent state and minimizes its squared distance to the goal-image latent state using CEM.CEM ranks candidates by this latent cost alone.Source: §2.1 · p. 3

2.2 COARSE READOUT OF THE SUCCESS-CRITERION QUANTITIES

  • On stack, baseline LeWM hand-position readout has 3.2 cm median error against a 2 cm tolerance; hand supervision reduces it to 0.14 cm.The three frozen pretrained encoders have median errors of 3.10 cm, 2.10 cm, and 4.02 cm, all outside tolerance.Source: §2.2 · p. 3

2.3 AUXILIARY LOSS SUPERVISED BY THE SUCCESS-CRITERION QUANTITIES

  • For each criterion quantity, a linear head reads both encoder and predictor outputs, and their squared errors against demonstration targets are added to the world-model loss.The encoder term averages over clip frames and the predictor term over one-step predictions; the auxiliary terms leave the original world-model loss unchanged.Source: §2.3 · p. 3
  • The shared head makes criterion quantities readable from both states, while discarding it after training preserves baseline planning cost, computation, and required inputs.The same head is needed because planning compares predictor-produced states with the goal-image latent state; training targets are quantities already recorded in demonstrations.Source: §2.3 · pp. 3–4

3 EXPERIMENTS

Across four manipulation tasks, supervising success-criterion quantities significantly improved success rate, though which quantities helped differed by task. Trial-stratum analysis showed that gains in one stratum could accompany losses in another.

3.1 EXPERIMENTAL SETUP

  • Evaluation paired 2,000 randomly drawn demonstration trials across baseline and auxiliary-loss models; each goal was 25 steps after the start, and trials could run up to 50 action steps.The trial set was shared between models, and a trial succeeded if its criterion held at any step within the action budget.Source: §3.1 · p. 4
  • The five tasks use goal tolerances on physical quantities; stack and coffee criteria were defined for this evaluation, with object tolerance 4 cm and hand tolerance 2 cm.A 4 cm hand tolerance made many trials successes at their start, so the hand tolerance was set to 2 cm; both tolerances were fixed before auxiliary-loss experiments.Source: §3.1 · p. 4, Table 1 · p. 5
  • The study compared baseline and auxiliary-loss models within scratch-trained LeWM and a frozen-DINOv2, predictor-only system, using paired two-sided McNemar tests rather than cross-system success-rate comparisons.Table 2 reports identical pairs of 2,000 trials for LeWM and 1,000 for the DINO-WM-type system; each model used one training seed.Source: §3.1 · pp. 4–5, Table 2 · p. 6

3.2 SUCCESS RATE BY SUPERVISED QUANTITY

  • In each of four manipulation tasks, some success-criterion supervision significantly improved success rate by 3–5% absolute, with the improving quantities varying by task.In stack, cube, and PushT, the improving model supervised all criterion quantities; in coffee, object-only supervision improved while supervising both did not.Source: §3.2 · p. 5, Table 2 · p. 6
  • Supervision of quantities outside a task’s success criterion showed no significant overall difference, and unwrapped Reacher joint angles also showed no effect because the same pose can have values 2π apart.Among in-criterion cases, four also left the overall rate unchanged: object supervision in PushT, hand supervision in coffee, and either quantity in stack.Source: §3.2 · p. 6, Table 2 · p. 6

3.3 EFFECT PER TRIAL STRATUM

  • On stack with LeWM, hand-only supervision improved reach-only trials but degraded manipulation trials, while object-only supervision had the reverse pattern.The authors call this complementary degradation; across the trials the baseline also failed, hand-only supervision brought the hand about 2 cm closer but left the cube about 10 cm short, while object-only supervision left the hand outside the 2 cm tolerance in reach-only trials.Source: §3.3 · pp. 6–7
  • In stack with LeWM, opposite stratum effects canceled in the overall rate; with frozen DINOv2-S, hand supervision improved both strata and object supervision strongly degraded reach-only trials.The frozen-system baseline succeeded in 5.7% of manipulation trials, so reach-only changes largely determined its overall rate.Source: §3.3 · p. 7
  • Across cube, PushT, stack, and coffee, models supervised with every criterion quantity did not significantly degrade either trial stratum.The stack model improved reach-only trials without degrading manipulation trials; the negative coffee manipulation-trial difference was not a significant degradation.Source: §3.3 · p. 7

4 RELATED WORK

Prior latent world models variously take physical quantities as inputs, reconstruct pixels, or omit such information, while representation-shaping and learned-cost methods pursue different targets. This work distinguishes itself by targeting criterion quantities for latent retention while leaving the planning cost unchanged.

Planning with latent world models

  • Prior world models supply physical information in different ways, but none asks whether criterion quantities are readable from the latent state at criterion precision; the success criterion, rather than the designer, determines the targets.Some systems concatenate proprioception, some reconstruct pixels, and LeWM omits the quantities; other training-time information is selected by the designer.Source: §4 · pp. 7–8
  • The intervention acts during training, unlike work that detects failures from latent states at test time [Zhang et al., 2026].Source: §4 · p. 9

Shaping representations with auxiliary losses

  • Auxiliary representation-shaping methods supervise different targets, and prior work does not test whether the latent distance used as planning cost reflects the success criterion.RC-aux supervises reachability [Li et al., 2026b], while PhyLatent and PSG-JEPA supervise physical quantities without contrasting in-criterion and out-of-criterion quantities or measuring effects by trial stratum [Zeng et al., 2026; Yan et al., 2026].Source: §4 · p. 9

Learning the distance used as a cost

  • Unlike methods that learn a planning distance from temporal proximity, transition costs, or executed outcomes, this approach leaves the cost unchanged and targets criterion quantities in the latent state.The cited learned-distance approaches include VIP and LIV [Ma et al., 2023b;a], Quasimetric RL, and TRM [Li et al., 2026a].Source: §4 · p. 9

5 CONCLUSION AND LIMITATIONS

Supervising every quantity in a task’s success criterion improved success across the eligible tasks without significant overall or per-stratum degradation. The authors recommend using the criterion as the training target, while noting limits on evidence and unresolved mechanisms.

Conclusion

  • All-criterion supervision improved success in all four eligible tasks and never significantly degraded it overall or by trial stratum; out-of-criterion supervision left success unchanged.The improvement also replicated on stack with a frozen pretrained encoder.Source: §5 · p. 9
  • Partial supervision helped some trials but hurt others on stack, while a corrected readout alone did not raise success.Supervising all criterion quantities is the authors’ practical rule and adds nothing to the model at test time.Source: §5 · p. 9

Limitations

  • Comparisons use a single training seed; the auxiliary-loss weight was swept only for hand in stack, and frozen-encoder replication was limited to stack and failed on frozen V-JEPA 2.The all-quantity stack models’ improvements ranged from +1.35% to +11.35%.Source: §5 · p. 9
  • For predicate-based criteria, the target to supervise is undetermined; on real robots, object positions may require motion capture or estimators.The route from the auxiliary loss to higher success rates and the harm from supervising only some quantities remain unexplained.Source: §5 · pp. 9–10

A IMPLEMENTATION DETAILS

The implementation combines a compact latent world model and CEM planning with linear auxiliary readouts, squared-error supervision, and task-specific training data and loss weights.

Model and planning

  • LeWM uses a 192-dimensional latent state and plans five predictor applications ahead, executing 25 steps before replanning.Its ViT-Tiny encoder has 5.5M parameters, and its causal Transformer predictor has 10.8M; CEM uses 300 candidates and 30 iterations.Source: A IMPLEMENTATION DETAILS · p. 13

Linear head and squared error

  • A shared linear head maps encoder and predictor latents to each target with squared error; targets are standardized, and criterion tolerances do not enter the loss.The linear readout is chosen because planning cost is Euclidean distance in latent space.Source: A IMPLEMENTATION DETAILS · p. 13

Auxiliary-loss weights and added parameters

  • Auxiliary heads add parameters, with λ = 1 for DINO-WM-type and frozen V-JEPA 2 systems; the hand weight in stack was tested at 1, 0.1, and 0.03.Other tasks and systems used default weights or a value tried on stack.Source: A IMPLEMENTATION DETAILS · p. 13

Tasks and training data

  • Tasks use 224×224 RGB observations and demonstration datasets; stack and coffee demonstrations were synthesized by MimicGen from a few human demonstrations.Reacher’s accumulated joint angles create targets 2π apart; a linear head’s median readout error was 0.53 rad, ten times the 0.05-rad tolerance.Source: A IMPLEMENTATION DETAILS · p. 13

B FROZEN-ENCODER SYSTEMS: SCOPE OF REPLICATION AND AN EXCEPTION

Frozen-encoder replication was limited to stack, and the frozen V-JEPA 2 system is an exception: improved hand readout accompanied lower success, which cost-side changes did not restore.

Scope of replication

The exception of the frozen V-JEPA 2

  • On stack, frozen V-JEPA 2 hand readout error fell from 1.18 cm to 0.58 cm, while success fell from 29.50% to 24.40% (−5.10%; n = 1000, p = 9 × 10−4).Replication on frozen-encoder systems was limited to stack.Source: B FROZEN-ENCODER SYSTEMS: SCOPE OF… · pp. 13–14

Interventions on the cost side

  • Whitening and reweighting latent distance along the hand-readout directions both failed to restore frozen V-JEPA 2 success to baseline.The cause of the success-rate drop remained unresolved despite improved readout.Source: B FROZEN-ENCODER SYSTEMS: SCOPE OF… · p. 14

C VALIDITY OF THE SUCCESS CRITERIA AND RESOLUTION OF THE TEST

The authors test tolerance sensitivity and paired-test resolution, and define trial strata by whether object displacement is required to meet the goal criterion.

Resolution of the test

  • With paired tests on identical trials, the significance threshold corresponds to 0.3% for 2000 trial pairs and 0.6% for 1000; these are test-imposed lower bounds, not measurement precision.Significance requires at least six disagreements, all favoring one model.Source: C VALIDITY OF THE SUCCESS CRITERIA AND… · pp. 14–15

Validity of the success criteria

  • Both main-text significant models remained significant across object tolerances from 4 cm to 1 cm; stack remained significant at hand tolerances of 3–4 cm, and coffee through 3 cm.At 1-cm hand tolerance, the baseline succeeded in only 0.7% of stack trials and 4.7% of coffee trials.Source: C VALIDITY OF THE SUCCESS CRITERIA AND… · p. 15, Table 6 · p. 15
  • The false-match rate was 0.15% in stack and 4.7% in coffee; in coffee, the pod condition alone accepted 24.6% of states from other episodes, while adding the hand condition reduced false matches.The paper attributes the pod-only rate presumably to the pod ending in a similar position across episodes.Source: C VALIDITY OF THE SUCCESS CRITERIA AND… · p. 15

Trial strata

  • A manipulation trial starts with object displacement at least the object tolerance; Figure 3 contrasts its start and goal frames with reach-only trials.Reacher has no object, while cube’s criterion has no hand term, so its reach-only trials already satisfy the criterion at the start.Source: C VALIDITY OF THE SUCCESS CRITERIA AND… · p. 15, Figure 3 · p. 16

D SUPPORTING EXPERIMENTS

The supporting experiments show that auxiliary-loss effects depend on which quantities are supervised and which trial strata test them. Readout, candidate ranking, criterion components, and end states help characterize those effects, but improvements in proxies do not consistently track success.

Readout of supervised and unsupervised quantities

  • Supervising one quantity also changed readout of quantities not supervised: hand supervision improved the carried-cube readout, while object supervision improved the carried-hand readout.For the untouched base cube, hand-only supervision raised readout error from 0.96 cm to 3.49 cm and frames beyond tolerance from 3% to 43%; with object supervision, terminal hand error fell from 1.47 cm to 1.13 cm and beyond-tolerance share from 26% to 15%.Source: D SUPPORTING EXPERIMENTS · pp. 16–17, Table 7 · p. 17

Ranking of candidates and shape of the search

  • Demo top-5 rate rose from 50% for baseline to 85% for hand-only λ = 1, yet that model's success-rate difference was −1.30%, the lowest of the compared models.The hand-only λ = 1 model had the best candidate ranking and most focused search, but ranking and search changes did not identify how auxiliary loss raised success.Source: D SUPPORTING EXPERIMENTS · p. 17, Figure 4 · p. 18

Component decomposition of PushT

  • On PushT, the hand-supervised model significantly improved only the agent component, while the unsupervised block component changed by +0.1%.The agent contributed 86% of the baseline criterion's median squared norm and exceeded the block term in 83% of trials; baseline block readout error was already inside the 15-unit component threshold.Source: D SUPPORTING EXPERIMENTS · p. 17, Table 8 · p. 18

Failure types in coffee and stack

  • Coffee baseline failures were 12% stop-short, 44% depart-after-approaching, and 44% never-approach, unlike stack failures, which were 57%, 15%, and 28%, respectively.The departing and never-approach types made up 88% of coffee failures; hand supervision brought the executed end state 1.4 cm closer but remained outside the 2 cm tolerance.Source: D SUPPORTING EXPERIMENTS · pp. 18–19

End states of the degraded trials

  • On manipulation trials where both models failed, hand-only supervision brought the hand closer to the goal (3.19 cm versus 5.07 cm) but left the object farther away (11.07 cm versus 9.83 cm).The object entered tolerance along the path in 15% versus 27% of trials; on reach-only trials, object-only supervision left the hand farther from goal (5.05 cm versus 4.36 cm).Source: D SUPPORTING EXPERIMENTS · p. 17, Table 9 · p. 19

Effects by trial stratum

  • Hand-only models improved reach-only trials by +14.8%, +14.5%, and +5.9% at λ = 0.03, 0.1, and 1, while manipulation trials degraded by −3.9%, −7.0%, and −8.2%, respectively.All five both-supervised models improved significantly on reach-only trials and none significantly degraded manipulation trials; the three λ = 0.03 seeds had overall gains of +11.4%, +3.7%, and +1.4%.Source: D SUPPORTING EXPERIMENTS · p. 19

Demanded object displacement and the effect

  • The hand-only effect was largest for 0.1–4 cm object displacement, declined as displacement grew, and became a significant degradation above 13 cm.Object-only supervision showed a mirror pattern, while both-quantity models improved at small displacement and did not degrade at large displacement.Source: D SUPPORTING EXPERIMENTS · p. 19

E ITEM-BY-ITEM COMPARISON WITH THE CLOSEST WORK

The comparison positions this work by its criterion-fixed supervision targets and its evaluation of effects by quantity and trial stratum. The accompanying figures describe stack results across trial strata and demanded object displacement.

Item-by-item comparison with the closest work

  • Only this work fixes auxiliary-loss targets to quantities in the success criterion and measures effects separately by supervised quantity and trial stratum.The comparison also distinguishes whether methods use latent-world-model planning, apply auxiliary losses to predicted latents, and leave inference unchanged.Source: E ITEM-BY-ITEM COMPARISON WITH THE… · p. 19, Table 10 · p. 20

Effects by trial stratum

  • Figure 5 compares success-rate differences from baseline across reach-only and manipulation strata for eight stack auxiliary-loss models, using paired trials and training seeds.The caption identifies 980 reach-only and 1020 manipulation trials and reports paired-bootstrap intervals and paired McNemar significance markers.Source: Figure 5 · p. 20

Demanded object displacement

  • Figure 6 bins stack trials by the larger cube's start-to-goal displacement and plots success-rate differences from baseline for four auxiliary-loss models.The five bins contain 645, 335, 236, 449, and 335 trials.Source: Figure 6 · p. 21
Loading 2610.01224v1…