Track manipulation, locomotion and embodied systems.
Search papers, labs, and topics across Lattice.
Latent world models struggle to encode task-critical success-criterion quantities (e.g., end-effector position) with sufficient precision for effective planning, leading to suboptimal performance. This work introduces a novel auxiliary loss that directly trains a linear head to regress these success-criterion quantities from the latent state, thereby improving the latent representation's fidelity to task success. This criterion-aligned auxiliary loss alone yields statistically significant improvements in success rates on benchmark tasks like PushT and cube.
Latent world models can be significantly improved by directly supervising their latent states with task success criteria, boosting performance without altering inference-time computation.
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.