Better world-model planning by training on success criteria

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

Machine LearningRobotics

Summary

Planning with world models involves guessing which actions will succeed based on internal representations called latent states. However, these latent states often do not accurately capture the important physical details that determine success, such as the position of a robot's hand. The authors show that by training the model to predict these exact success-related physical quantities during learning, the model can better distinguish successful plans from failures. This extra training step improves success rates on robotic tasks without changing how the model works during actual planning. It highlights the importance of aligning training goals with what actually matters for success.

What this means in practice

  • For robotics engineers: Improve robotic manipulation success by training models to explicitly represent key physical success metrics in latent states.
  • For autonomous system developers: Enhance planning accuracy in autonomous agents by adding auxiliary training targets aligned with task success criteria.

Authors

Takumi Hara, Kanata Suzuki

Abstract

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.
View PDFOpen arXiv