Completion aware guidance improves robot action predictions for tasks

Completion Aware Guidance for World Action Models

RoboticsArtificial IntelligenceMachine Learning

Summary

World Action Models try to predict what a robot will do next by imagining possible futures, but sometimes these guesses miss the key step needed to finish a task even if they look right locally. The authors found this problem happens when the models focus on short-term moves rather than the full task. They introduced a new technique called Completion Aware Guidance that helps guide the robot’s imagined actions toward actually finishing the task. This method improved task success rates and reduced incomplete predictions without needing to retrain the models.

What this means in practice

  • For robot developers: Enhance robot control systems to better predict and complete tasks by reducing incomplete imagined actions in robot planning.
  • For simulation engineers: Improve simulation models used to test robot actions by increasing success rates and decreasing task-incomplete predictions through guided sampling.

Authors

Seungyeon Kim, Junhoo Lee, Baekseung Kim, Minkyu Kim, Nojun Kwak

Abstract

World Action Models (WAMs) predict visual futures and robot actions, yet they remain susceptible to task-incomplete imagination, where plausible, action-consistent predictions omit the transition needed for task completion. In this paper, we show that this failure is not inherent to the world model backbone, but emerges when adapted for short-chunk control, which can repeatedly favor plausible local continuations over task-completing transitions. To address this, we introduce Completion Aware Guidance (CAG), a training-free sampling method that guides generation toward task completion. Across representative WAMs, CAG improves success from 64% to 70% on a RoboTwin 2.0 subset and from 69% to 75% in zero-shot simulation, while reducing task-incomplete imagination from 79% to 40%.