Visuomotor policies improve task success with real-time correction
Back on Track: Deployment-time Policy Improvement via World Model Planning
Robotics
Summary
Visuomotor policies help robots see and act to complete tasks, but sometimes they get stuck or fail without knowing why. The authors create a system that watches the robot’s actions and notices when it's not making progress. When that happens, it uses a kind of 'world model' to quickly plan a fix and steer the robot back on track before letting it continue. This approach works with existing robot policies and makes them better at finishing tasks without needing extra retraining.
What this means in practice
- •For robotics engineers: Enhance robot manipulators to detect and correct execution errors without retraining during deployment.
- •For automation system integrators: Integrate monitoring and planning layers over existing black-box policies to improve operational robustness in industrial automation.
Authors
Colton Crosby, Hao Wang, Mac Schwager, Somil Bansal
Abstract
State-of-the-art visuomotor policies are increasingly capable of executing complex tasks, yet can fail during deployment in ways they cannot predict or recover from on their own. Improving their performance and robustness typically requires collecting more data, retraining, or modifying the policy architecture, all of which are time- and resource-intensive in practice. We instead propose an offline-trained monitoring-and-intervention framework that improves a policy during deployment using only black-box query access. From a limited set of labeled on-policy rollouts, we learn a policy-conditioned liveness value that measures the nominal policy's progress toward task completion; we additionally train a compact latent world model using these rollouts together with play data. At deployment, the liveness value continuously monitors execution. When it detects a loss of progress, a short-horizon planner searches the world model for corrective actions using the same value as its planning objective, after which control returns to the nominal policy. We evaluate our method on the LIBERO-Goal task suite and on a Franka Panda manipulator, using the vision-language-action model $π_{0.5}$ and a diffusion policy as nominal policies, and demonstrate significant improvements in task success rates, while supporting real-time control. Additional results and videos are available on the project website (https://policy-improvement-wm.pages.dev/).