Direct prediction method speeds up robot control with high task success
De-Diffusing WAMs: Action-Sufficient Direct Prediction for Robot Control
RoboticsArtificial Intelligence
Summary
Robots often use complex models that predict future visual scenes to decide their next moves, but these can be slow because they rely on many repeated steps. The authors studied what parts of the predicted future really matter for making accurate actions in just one step. They developed DirectWAM, a way to teach a robot to predict future states that are just good enough for immediate action decisions, which makes the robot act faster without losing accuracy. This new method was tested on various simulated and real tasks, showing it works well and runs much faster than older approaches.
What this means in practice
- •For robotics engineers: Implement faster control systems by replacing iterative future prediction with direct single-step prediction that preserves action accuracy.
- •For automation system developers: Build more responsive manipulation or interaction tasks by using distilled models that reduce inference time without reducing task success rates.
Authors
Shaoyi Zheng, Zhongtao Jiang, Yue Peng, Wenbo Lu, Shunchi Zhang, Xi Wang, Shengji Wan
Abstract
Cascade world-action models (WAMs) use predicted visual futures to condition robot actions, but inherit costly iterative inference from diffusion-pretrained experts. We ask what a single-step predicted future must preserve to support accurate single-step action prediction. In the evaluated WAM, resampling video noise has little effect on downstream actions, yet naively truncating both streams to one step substantially reduces task success. Video-space supervision only partially closes this gap, and a local sensitivity analysis shows that vision prediction errors have highly unequal consequences for action accuracy. These findings motivate DirectWAM, which converts pretrained video and action experts into single-step predictors without stochastic sampling through action-sufficient distillation, while retaining the explicit future-to-action cascade. After learning a one-step action model, we train the future predictor by decoding actions from its own predicted futures and backpropagating demonstration action error through the expert, with an optional future-latent reconstruction loss. The resulting objective evaluates the student-generated future through its downstream control consequences rather than relying on representation matching alone. DirectWAM achieves 97.1% average success across four LIBERO suites and 85.4% on RoboTwin 2.0, outperforming Flash-WAM at the same one-evaluation-per-stream budget. This advantage has also extended to three real-world tasks that we designed. Compared with the full iterative teacher, it reduces inference latency by $28.0\times$.